model fusion and onnx export

models/onnx_export.py

@@ -1,6 +1,6 @@
-# Exports a pytorch *.pt model to *.onnx format. Example usage (run from ./yolov5 directory):
-# $ export PYTHONPATH="$PWD"
-# $ python models/onnx_export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
+# Exports a pytorch *.pt model to *.onnx format
+# Example usage (run from ./yolov5 directory):
+# $ export PYTHONPATH="$PWD" && python models/onnx_export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
 
 import argparse
 
@@ -10,10 +10,11 @@ from models.common import *
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', default='./weights/yolov5s.pt', help='weights path')
-    parser.add_argument('--img-size', default=640, help='inference size (pixels)')
-    parser.add_argument('--batch-size', default=1, help='batch size')
+    parser.add_argument('--weights', type=str, default='./weights/yolov5s.pt', help='weights path')
+    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
+    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
     opt = parser.parse_args()
+    print(opt)
 
     # Parameters
     f = opt.weights.replace('.pt', '.onnx')  # onnx filename
@@ -23,7 +24,7 @@ if __name__ == '__main__':
     google_utils.attempt_download(opt.weights)
     model = torch.load(opt.weights)['model']
     model.eval()
-    # model.fuse()  # optionally fuse Conv2d + BatchNorm2d layers TODO
+    # model.fuse()
 
     # Export to onnx
     model.model[-1].export = True  # set Detect() layer export=True