feat(scripts): 数据转换、ONNX 导出与推理性能分析

- convert_annotations: labelme / VOC -> YOLO txt，顺带生成 data.yaml - split_dataset: 切分成 images/{train,val} 与 labels/{train,val}，默认建软链不复制数据 - export_onnx: 导出部署用 ONNX，CPU 可跑，缺依赖时提示 `uv sync --group export` - profile_inference: torch.profiler 测单帧延迟/FPS、CUDA 显存峰值与算子级明细 - check_env: Python / PyTorch / Nvidia GPU 自检 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-02 17:55:20 +08:00
parent b0d0dfffbb
commit 4b1147ed8f
5 changed files with 561 additions and 0 deletions
@@ -0,0 +1,77 @@
+#!/usr/bin/env python
+"""导出 ONNX，供实时设备（TensorRT / OpenVINO）部署。
+
+    python scripts/export_onnx.py --config configs/train/baseline.yaml \
+        --checkpoint runs/baseline/best.pth --output runs/baseline/model.onnx
+
+导出在 CPU 上完成，不需要 GPU；需要额外依赖：
+    uv sync --group export
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+import torch  # noqa: E402
+
+from datasets import load_classes  # noqa: E402
+from engine import load_checkpoint  # noqa: E402
+from models import build_model  # noqa: E402
+from utils import load_config  # noqa: E402
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--opset", type=int, default=17)
+    parser.add_argument("--dynamic-batch", action="store_true")
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    model_cfg = dict(cfg["model"])
+    model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"])))
+    model_cfg.setdefault("label_mode", str(cfg.get("label_mode", "multilabel")))
+    model = build_model(model_cfg)
+    load_checkpoint(args.checkpoint, model, map_location="cpu")
+    model.eval()
+
+    height, width = cfg.get("image_size", [512, 512])
+    channels = 1 if cfg.get("gray") else 3
+    dummy = torch.randn(1, channels, int(height), int(width))
+
+    output_names = ["heatmap", "wh", "offset"] + (["logits"] if model.cls_head is not None else [])
+    dynamic_axes = {"images": {0: "batch"}} if args.dynamic_batch else None
+    if dynamic_axes:
+        dynamic_axes.update({name: {0: "batch"} for name in output_names})
+
+    try:
+        import onnxscript  # noqa: F401
+    except ImportError:
+        print("缺少导出依赖，请先执行: uv sync --group export", file=sys.stderr)
+        return 1
+
+    out_path = Path(args.output)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    torch.onnx.export(
+        model,
+        dummy,
+        str(out_path),
+        input_names=["images"],
+        output_names=output_names,
+        opset_version=args.opset,
+        dynamic_axes=dynamic_axes,
+    )
+    print(f"已导出 {out_path}（输入 {tuple(dummy.shape)}，输出 stride={model.out_stride}）")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())