Shaun ZhouandClaude Opus 5
4b1147ed8f
- 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>
78 lines
2.5 KiB
Python
Executable File
78 lines
2.5 KiB
Python
Executable File
#!/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())
|