Dynamic batchsize export support for ONNX and TensorRT

README.md

@@ -153,26 +153,35 @@ python detect.py --weights yolov7.pt --conf 0.25 --img-size 640 --source inferen
 
 
 ## Export
-
-Pytorch -> ONNX -> TensorRT -> Detection on TensorRT in Python <a href="https://colab.research.google.com/gist/AlexeyAB/fcb47ae544cf284eb24d8ad8e880d45c/yolov7trtlinaom.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
-
-
-**Pytorch to ONNX**, use `--include-nms` flag for the end-to-end ONNX model with `EfficientNMS`
+### Pytorch to ONNX
+Tested with: Python 3.7.13 and Pytorch 1.12.0+cu113
+
+- `--include-grid`: Add Detect layer
+- `--include-nms`: Add EfficientNMS plugin
+- `--include-nms-score-thresh`: NMS plugin threshold for score
+- `--include-nms-nms-thresh`: NMS plugin threshold for NMS IoU
+- `--include-nms-detections-per-image`: NMS plugin threshold for maximum amount of resulting objects per image
+- `--dynamic-batch`: Make first input dimension dynamic for dynamic batch support
+- `--dynamic-shape`: Make third and fourth input dimension dynamic for dynamic width and height
+- `--onnx-simplify` Run [onnx-simplifier](https://github.com/daquexian/onnx-simplifier)
+
+yolov7-tiny to ONNX with dynamic batchsize, grid and nms(score-thresh=0.2, nms-thresh=0.5, detections-per-image=500):
 ```shell
 wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt
-python export.py --weights yolov7-tiny.pt --grid --include-nms
+!python export.py --weights yolov7-tiny.pt --onnx-simplify --dynamic-batch --include-grid --include-nms --include-nms-score-thresh=0.2 --include-nms-nms-thresh=0.5 --include-nms-detections-per-image=500
 ```
 
-**ONNX to TensorRT**
+### ONNX to TensorRT with docker
 ```shell
-git clone https://github.com/Linaom1214/tensorrt-python.git
-cd tensorrt-python
-python export.py -o yolov7-tiny.onnx -e yolov7-tiny-nms.trt -p fp16
+$ docker run -it --rm --gpus=all nvcr.io/nvidia/tensorrt:22.04-py3
+$ # from new shell copy onnx to container
+$ docker cp yolov7-tiny.onnx 898c16f38c99:/workspace/tensorrt/bin
+$ # in container now
+$ cd /workspace/tensorrt/bin
+$ # convert onnx to tensorrt with min batch size 1, opt batch size 8 and max batch size 16
+$ ./trtexec --onnx=yolov7-tiny.onnx --minShapes=input:1x3x640x640 --optShapes=input:8x3x640x640 --maxShapes=input:16x3x640x640 --fp16 --workspace=4096 --saveEngine=yolov7-tiny.engine
 ```
 
-Tested with: Python 3.7.13, Pytorch 1.12.0+cu113
-
-
 ## Citation
 
 ```

export.py

@@ -19,11 +19,15 @@
     parser.add_argument('--weights', type=str, default='./yolor-csp-c.pt', help='weights path')
     parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')  # height, width
     parser.add_argument('--batch-size', type=int, default=1, help='batch size')
-    parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
-    parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
+    parser.add_argument('--dynamic-batch', action='store_true', help='dynamic ONNX batchsize')
+    parser.add_argument('--dynamic-shape', action='store_true', help='dynamic ONNX input width and height')
     parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
-    parser.add_argument('--simplify', action='store_true', help='simplify onnx model')
-    parser.add_argument('--include-nms', action='store_true', help='export end2end onnx')
+    parser.add_argument('--include-grid', action='store_true', help='export Detect() layer grid')
+    parser.add_argument('--include-nms', action='store_true', help='export nms plugin')
+    parser.add_argument('--include-nms-score-thresh', default=0.25, type=float, help='export nms plugin score threshold, default 0.25')
+    parser.add_argument('--include-nms-nms-thresh', default=0.45, type=float, help='export nms plugin nms threshold, default 0.45')
+    parser.add_argument('--include-nms-detections-per-image', default=100, type=int, help='export nms plugin max detections per image, default 100')
+    parser.add_argument('--onnx-simplify', action='store_true', help='simplify onnx model')
     opt = parser.parse_args()
     opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
     print(opt)
@@ -52,7 +56,7 @@
                 m.act = SiLU()
         # elif isinstance(m, models.yolo.Detect):
         #     m.forward = m.forward_export  # assign forward (optional)
-    model.model[-1].export = not opt.grid  # set Detect() layer grid export
+    model.model[-1].export = not opt.include_grid  # set Detect() layer grid export
     y = model(img)  # dry run
     if opt.include_nms:
         model.model[-1].include_nms = True
@@ -74,10 +78,21 @@
         print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
         f = opt.weights.replace('.pt', '.onnx')  # filename
         model.eval()
-        torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
+        dynamic_axes = None
+        if opt.dynamic_batch or opt.dynamic_shape:
+            dynamic_axes = { 'input': {}, 'output': {} }
+            if opt.dynamic_batch:
+                dynamic_axes['input'][0] = 'batch'
+                dynamic_axes['output'][0] = 'batch'
+            if opt.dynamic_shape:
+                dynamic_axes['input'][2] = 'height'
+                dynamic_axes['input'][3] = 'width'
+                dynamic_axes['output'][2] = 'y'
+                dynamic_axes['output'][3] = 'x'
+        torch.onnx.export(model, img, f, verbose=False, opset_version=12,
+                          input_names=['input'],
                           output_names=['classes', 'boxes'] if y is None else ['output'],
-                          dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # size(1,3,640,640)
-                                        'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
+                          dynamic_axes=dynamic_axes)
 
         # Checks
         onnx_model = onnx.load(f)  # load onnx model
@@ -91,7 +106,7 @@
         #     meta.key, meta.value = k, str(v)
         # onnx.save(onnx_model, f)
 
-        if opt.simplify:
+        if opt.onnx_simplify:
             try:
                 import onnxsim
 
@@ -106,7 +121,7 @@
         if opt.include_nms:
             print('Registering NMS plugin for ONNX...')
             mo = RegisterNMS(f)
-            mo.register_nms()
+            mo.register_nms(score_thresh=opt.include_nms_score_thresh, nms_thresh=opt.include_nms_nms_thresh, detections_per_img=opt.include_nms_detections_per_image)
             mo.save(f)
 
     except Exception as e: