luisarizmendi commited on
Commit
9a77e24
·
1 Parent(s): 0f86375

pt-to-onnx script

Browse files
Files changed (1) hide show
  1. tools/pytorch-to-onnx.py +13 -0
tools/pytorch-to-onnx.py ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ from ultralytics import YOLO
3
+
4
+ # Parse command-line arguments
5
+ parser = argparse.ArgumentParser(description="Export a YOLO model to ONNX format.")
6
+ parser.add_argument("model_path", type=str, help="Path to the YOLO model (.pt file)")
7
+ args = parser.parse_args()
8
+
9
+ # Load and export the model
10
+ model = YOLO(args.model_path)
11
+ model.export(format='onnx', imgsz=640)
12
+
13
+