Update pytorch2onnx.py

Add the exporting graph as dynamics graph. In Jetson Xavier JetPack 4.4, multi-batch engine could not be built with statics engine.

Author: KelvinCPChiu

tool/pytorch2onnx.py

@@ -2,7 +2,7 @@
 import torch
 
 
-def transform_to_onnx(cfgfile, weightfile, batch_size=1):
+def transform_to_onnx(cfgfile, weightfile, batch_size=1, dynamics=False):
     model = Darknet(cfgfile)
 
     model.print_network()
@@ -13,20 +13,39 @@ 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,
-                      x,
-                      onnx_file_name,
-                      export_params=True,
-                      opset_version=11,
-                      do_constant_folding=True,
-                      dynamic_axes=None)
+        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
+        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,
+                          opset_version=11,
+                          do_constant_folding=True,
+                          dynamic_axes=None)
+
+        print('Onnx model exporting done')
+        return onnx_file_name
 
 
 if __name__ == '__main__':
@@ -39,6 +58,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])
+        dynamics = True if sys.argv[4] == 'True' else False
+        transform_to_onnx(cfgfile, weightfile, batch_size, dynamics)
     else:
         print('Please execute this script this way:\n')
         print('  python darknet2onnx.py <cfgFile> <weightFile>')