Adding ONNX info + fix for export

README.md

@@ -64,6 +64,16 @@ python train.py --data coco2017.data --cfg yolov4-pacsp.cfg --weights '' --name
 python test_half.py --data coco2017.data --cfg yolov4-pacsp.cfg --weights yolov4-pacsp.pt --img 736 --iou-thr 0.7 --batch-size 8
 ```
 
+## ONNX export
+
+Set the [`ONNX_EXPORT` variable to `True` in models.py](https://github.com/WongKinYiu/PyTorch_YOLOv4/blob/master/models.py#L5), then launch the detection script with the desired model
+
+```
+python detect.py --cfg yolov4-pacsp-s.cfg --weights weights/yolov4-pacsp-s.pt
+```
+
+The ONNX input size is [hardcoded in detect.py](https://github.com/WongKinYiu/PyTorch_YOLOv4/blob/master/detect.py#L10) when `ONNX_EXPORT` is set to true.
+
 ## Citation
 
 ```

detect.py

@@ -5,9 +5,8 @@
 from utils.datasets import *
 from utils.utils import *
 
-
 def detect(save_img=False):
-    img_size = (320, 192) if ONNX_EXPORT else opt.img_size  # (320, 192) or (416, 256) or (608, 352) for (height, width)
+    img_size = (320, 320) if ONNX_EXPORT else opt.img_size  # (320, 192) or (416, 256) or (608, 352) for (height, width)
     out, source, weights, half, view_img, save_txt = opt.output, opt.source, opt.weights, opt.half, opt.view_img, opt.save_txt
     webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')
 
@@ -45,14 +44,15 @@ def detect(save_img=False):
         model.fuse()
         img = torch.zeros((1, 3) + img_size)  # (1, 3, 320, 192)
         f = opt.weights.replace(opt.weights.split('.')[-1], 'onnx')  # *.onnx filename
-        torch.onnx.export(model, img, f, verbose=False, opset_version=11,
-                          input_names=['images'], output_names=['classes', 'boxes'])
+        torch.onnx.export(model, img, f, verbose=True, opset_version=10, export_params=True, do_constant_folding=True,
+                          input_names=['images'], output_names=['classes', 'boxes'], dynamic_axes=None)
 
         # Validate exported model
         import onnx
         model = onnx.load(f)  # Load the ONNX model
         onnx.checker.check_model(model)  # Check that the IR is well formed
         print(onnx.helper.printable_graph(model.graph))  # Print a human readable representation of the graph
+        print("Exported to " + f)
         return
 
     # Half precision

models.py

@@ -14,6 +14,9 @@ def create_modules(module_defs, img_size, cfg):
     module_list = nn.ModuleList()
     routs = []  # list of layers which rout to deeper layers
     yolo_index = -1
+    # To avoid dimensions issues when exporting to ONNX
+    # https://github.com/ultralytics/yolov3/issues/1169#issuecomment-650678238
+    upsample_index = 0
 
     for i, mdef in enumerate(module_defs):
         modules = nn.Sequential()
@@ -70,7 +73,9 @@ def create_modules(module_defs, img_size, cfg):
 
         elif mdef['type'] == 'upsample':
             if ONNX_EXPORT:  # explicitly state size, avoid scale_factor
-                g = (yolo_index + 1) * 2 / 32  # gain
+                # g = (yolo_index + 1) * 2 / 32  # gain
+                g = (upsample_index + 1) * 2 / 32  # gain
+                upsample_index += 1
                 modules = nn.Upsample(size=tuple(int(x * g) for x in img_size))  # img_size = (320, 192)
             else:
                 modules = nn.Upsample(scale_factor=mdef['stride'])
@@ -98,9 +103,13 @@ def create_modules(module_defs, img_size, cfg):
 
         elif mdef['type'] == 'yolo':
             yolo_index += 1
-            stride = [8, 16, 32]  # P5, P4, P3 strides
-            if any(x in cfg for x in ['yolov4-tiny']):  # stride order reversed
-                stride = [32, 16, 8]
+            #stride = [8, 16, 32]  # P5, P4, P3 strides
+            #if any(x in cfg for x in ['yolov4-tiny']):  # stride order reversed
+            #    stride = [32, 16, 8]
+            stride = [32, 16, 8]  # P5, P4, P3 strides
+            if any(x in cfg for x in ['panet', 'yolov4', 'cd53']):  # stride order reversed
+                stride = list(reversed(stride))
+
             layers = mdef['from'] if 'from' in mdef else []
             modules = YOLOLayer(anchors=mdef['anchors'][mdef['mask']],  # anchor list
                                 nc=mdef['classes'],  # number of classes