Update darknet2onnx.py

Add the dynamics axes for multi-batch cases.

Author: KelvinCPChiu

tool/darknet2onnx.py

@@ -3,7 +3,7 @@
 from tool.darknet2pytorch import Darknet
 
 
-def transform_to_onnx(cfgfile, weightfile, batch_size=1):
+def transform_to_onnx(cfgfile, weightfile, batch_size=1,dynamic):
     model = Darknet(cfgfile)
 
     model.print_network()
@@ -14,11 +14,30 @@ def transform_to_onnx(cfgfile, weightfile, batch_size=1):
 
     x = torch.randn((batch_size, 3, model.height, model.width), requires_grad=True)  # .cuda()
 
-    onnx_file_name = "yolov4_{}_3_{}_{}.onnx".format(batch_size, model.height, model.width)
+    if dynamics:
+        
+        onnx_file_name = "yolov4_{}_3_{}_{}_dyna.onnx".format(batch_size, model.height, model.width)
+        input_names = ["input"]
+        output_names = ['boxes', 'confs']
 
-    # Export the model
-    print('Export the onnx model ...')
-    torch.onnx.export(model,
+        dynamic_axes = {"input": {0: "batch_size"}, "boxes": {0: "batch_size"}, "confs": {0: "batch_size"}}
+        # Export the model
+        print('Export the onnx model ...')
+        torch.onnx.export(model,
+                          x,
+                          onnx_file_name,
+                          export_params=True,
+                          opset_version=11,
+                          do_constant_folding=True,
+                          input_names=input_names, output_names=output_names,
+                          dynamic_axes=dynamic_axes)
+
+        print('Onnx model exporting done')
+        return onnx_file_name
+
+    else:
+        onnx_file_name = "yolov4_{}_3_{}_{}_static.onnx".format(batch_size, model.height, model.width)
+        torch.onnx.export(model,
                       x,
                       onnx_file_name,
                       export_params=True,
@@ -27,8 +46,9 @@ def transform_to_onnx(cfgfile, weightfile, batch_size=1):
                       input_names=['input'], output_names=['boxes', 'confs'],
                       dynamic_axes=None)
 
-    print('Onnx model exporting done')
-    return onnx_file_name
+        print('Onnx model exporting done')
+        return onnx_file_name
+  
 
 
 if __name__ == '__main__':
@@ -41,6 +61,12 @@ def transform_to_onnx(cfgfile, weightfile, batch_size=1):
         weightfile = sys.argv[2]
         batch_size = int(sys.argv[3])
         transform_to_onnx(cfgfile, weightfile, batch_size)
+    elif len(sys.argv) == 5:
+        cfgfile = sys.argv[1]
+        weightfile = sys.argv[2]
+        batch_size = int(sys.argv[3])
+        dynamic = True if sys.argv[4] == 'True' else False
+        transform_to_onnx(cfgfile, weightfile, batch_size, dynamic)
     else:
         print('Please execute this script this way:\n')
         print('  python darknet2onnx.py <cfgFile> <weightFile>')