Spaces:
Sleeping
Sleeping
File size: 2,110 Bytes
cabc6bd | 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 | """Export the trained model to edge-deployment formats (plan 3.2).
Requires the `export` dependency group. From repo root:
uv run --group export python scripts/export_model.py # onnx + openvino + coreml
uv run --group export python scripts/export_model.py --formats onnx
uv run --group export python scripts/export_model.py --formats engine # TensorRT — NVIDIA GPU only (Colab)
Artifacts land next to the weights (runs/detect/.../weights/) and are
git-ignored — re-run this script to regenerate them. Benchmark them against the
PyTorch baseline with scripts/bench_exports.py.
"""
from __future__ import annotations
import argparse
from pathlib import Path
from ultralytics import YOLO
WEIGHTS = "runs/detect/yolov8n_v1_train/weights/best.pt"
DEFAULT_FORMATS = ["onnx", "openvino", "coreml"]
def dir_size_mb(path: Path) -> float:
if path.is_file():
return path.stat().st_size / 1e6
return sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) / 1e6
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--weights", default=WEIGHTS)
p.add_argument(
"--formats",
default=",".join(DEFAULT_FORMATS),
help="comma-separated Ultralytics export formats (onnx, openvino, coreml, engine, tflite, ...)",
)
p.add_argument("--imgsz", type=int, default=640)
p.add_argument("--half", action="store_true", help="FP16 export (GPU formats like engine)")
args = p.parse_args()
exported: list[tuple[str, Path]] = []
for fmt in args.formats.split(","):
fmt = fmt.strip()
# Fresh model per format: export() mutates model state in some backends.
model = YOLO(args.weights)
print(f"\n=== Exporting {fmt} ===")
out = model.export(format=fmt, imgsz=args.imgsz, half=args.half)
exported.append((fmt, Path(out)))
print(f"\n{'Format':<10} {'Size (MB)':>10} Path")
print("-" * 60)
for fmt, path in exported:
print(f"{fmt:<10} {dir_size_mb(path):>10.1f} {path}")
if __name__ == "__main__":
main()
|