Merge pull request #193 from KelvinCPChiu/master

Update the Torch to ONNX for dynamic graph.

Author: Tianxiaomo

DeepStream/Readme.md

@@ -0,0 +1 @@
+

README.md

@@ -212,6 +212,32 @@ you can use darknet2pytorch to convert it yourself, or download my converted mod
     
     Note:Errors will occur when using "pip install onnx-tf", at least for me,it is recommended to use source code installation
 
+# 7. ONNX2TensorRT and DeepStream Inference
+  
+  1. Compile the DeepStream Nvinfer Plugin 
+  
+  ```
+      cd DeepStream
+      make 
+  ```
+  2. Build a TRT Engine.
+  
+   For single batch, 
+   ```
+   trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
+   ```
+   
+   For multi-batch, 
+  ```
+  trtexec --onnx=<onnx_file> --explicitBatch --shapes=input:Xx3xHxW --optShapes=input:Xx3xHxW --maxShapes=input:Xx3xHxW --minShape=input:1x3xHxW --saveEngine=<tensorRT_engine_file> --fp16
+  ```
+  
+  Note :The maxShapes could not be larger than model original shape.
+  
+  3. Write the deepstream config file for the TRT Engine.
+  
+  
+   
 Reference:
 - https://github.com/eriklindernoren/PyTorch-YOLOv3
 - https://github.com/marvis/pytorch-caffe-darknet-convert

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=False):
     model = Darknet(cfgfile)
 
     model.print_network()
@@ -14,21 +14,40 @@ 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 dynamic:
 
-    # 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'], output_names=['boxes', 'confs'],
-                      dynamic_axes=None)
+        onnx_file_name = "yolov4_{}_3_{}_{}_dyna.onnx".format(batch_size, model.height, model.width)
+        input_names = ["input"]
+        output_names = ['boxes', 'confs']
 
-    print('Onnx model exporting done')
-    return onnx_file_name
+        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,
+                          opset_version=11,
+                          do_constant_folding=True,
+                          input_names=['input'], output_names=['boxes', 'confs'],
+                          dynamic_axes=None)
+
+        print('Onnx model exporting done')
+        return onnx_file_name
 
 
 if __name__ == '__main__':
@@ -41,6 +60,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>')

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, dynamic=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])
+        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>')