File size: 5,252 Bytes
67c84d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | #!/usr/bin/env python3
# Reference: https://github.com/infracv/rf-detr-cpp/blob/develop/trt-files/scripts/export_onnx.py
"""Export an RF-DETR PyTorch checkpoint to ONNX, plus a meta-sidecar JSON.
The sidecar is consumed by `rfdetr/core/engine_meta.hpp`. It captures the
variant identity, input H/W, query count, normalization stats, and color
order — i.e. everything the C++ runtime needs that is NOT recoverable from the
engine's tensor shapes alone.
Example:
python export_onnx.py --variant small --out-dir onnx
python export_onnx.py --variant seg-large --weights /tmp/seg-large.pth
"""
from __future__ import annotations
import argparse
import json
import os
import shutil
import sys
from pathlib import Path
from typing import Any, Dict
VARIANT_TABLE: Dict[str, Dict[str, Any]] = {
"nano": {"resolution": 384, "num_queries": 300, "patch": 16, "has_masks": False,
"ctors": ("RFDETRNano",)},
"small": {"resolution": 512, "num_queries": 300, "patch": 16, "has_masks": False,
"ctors": ("RFDETRSmall",)},
"medium": {"resolution": 576, "num_queries": 300, "patch": 16, "has_masks": False,
"ctors": ("RFDETRMedium",)},
"base": {"resolution": 560, "num_queries": 300, "patch": 14, "has_masks": False,
"ctors": ("RFDETRBase",)},
"large": {"resolution": 704, "num_queries": 300, "patch": 16, "has_masks": False,
"ctors": ("RFDETRLarge",)},
"seg-nano": {"resolution": 312, "num_queries": 100, "patch": 12, "has_masks": True,
"ctors": ("RFDETRSegNano",)},
"seg-small": {"resolution": 384, "num_queries": 100, "patch": 12, "has_masks": True,
"ctors": ("RFDETRSegSmall",)},
"seg-medium": {"resolution": 432, "num_queries": 200, "patch": 12, "has_masks": True,
"ctors": ("RFDETRSegMedium",)},
"seg-large": {"resolution": 504, "num_queries": 200, "patch": 12, "has_masks": True,
"ctors": ("RFDETRSegLarge",)},
"seg-xlarge": {"resolution": 624, "num_queries": 300, "patch": 12, "has_masks": True,
"ctors": ("RFDETRSegXLarge",)},
"seg-2xlarge": {"resolution": 768, "num_queries": 300, "patch": 12, "has_masks": True,
"ctors": ("RFDETRSeg2XLarge",)},
"seg-preview": {"resolution": 432, "num_queries": 200, "patch": 12, "has_masks": True,
"ctors": ("RFDETRSegPreview",)},
}
def resolve_ctor(rfdetr_module, ctor_names):
for name in ctor_names:
if hasattr(rfdetr_module, name):
return getattr(rfdetr_module, name)
raise AttributeError(
f"none of {ctor_names} are exposed by `rfdetr` — upstream may have renamed the variant. "
f"Edit VARIANT_TABLE in this script to match."
)
def main():
ap = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--variant", required=True, choices=sorted(VARIANT_TABLE.keys()),
help="RF-DETR variant to export")
ap.add_argument("--weights", default=None,
help="path to .pth checkpoint (default: upstream pretrained)")
ap.add_argument("--out-dir", default="onnx",
help="output directory for .onnx and .json sidecar")
ap.add_argument("--name", default=None,
help="output basename (default: rf-detr-<variant>)")
ap.add_argument("--opset", type=int, default=17)
ap.add_argument("--dynamic-batch", action="store_true",
help="export with batch dimension as a dynamic axis")
args = ap.parse_args()
spec = VARIANT_TABLE[args.variant]
try:
import rfdetr # type: ignore
except ImportError:
sys.exit("rfdetr Python package not installed (pip install rfdetr)")
ctor = resolve_ctor(rfdetr, spec["ctors"])
model_kwargs = {"pretrain_weights": args.weights} if args.weights else {}
model = ctor(**model_kwargs)
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
basename = args.name or f"rf-detr-{args.variant}"
# rfdetr's `.export(format="onnx", output_dir=...)` typically writes
# `inference_model.onnx` (and possibly a sim'd variant) into output_dir.
# We invoke it, then locate the produced .onnx and rename to our convention.
print(f"[export_onnx] exporting variant={args.variant} to {out_dir}/")
model.export(
format="onnx",
output_dir=str(out_dir),
opset_version=args.opset,
dynamic_batch=args.dynamic_batch,
)
# Find the most recently produced .onnx in out_dir; rename to <basename>.onnx.
candidates = sorted(out_dir.glob("*.onnx"), key=lambda p: p.stat().st_mtime, reverse=True)
if not candidates:
sys.exit(f"export reported success but no .onnx file appeared in {out_dir}")
src = candidates[0]
dst = out_dir / f"{basename}.onnx"
if src.resolve() != dst.resolve():
shutil.move(str(src), str(dst))
print(f"[export_onnx] renamed {src.name} -> {dst.name}")
print(f"[export_onnx] wrote {dst}")
if __name__ == "__main__":
main() |