Repository: ultralytics/yolov5

Issue:
Exporting TensorRT (engine) with dynamic batch size failing
### Search before asking

- [X] I have searched the YOLOv5 [issues](https://github.com/ultralytics/yolov5/issues) and found no similar bug report.


### YOLOv5 Component

Export

### Bug

I'm trying to export the pretrained yolov5 to `engine` format with dynamic shape using this command -

`python export.py --weights yolov5m_coco.pt --include engine --device 0 --batch-size 256 --dynamic`

But it's failing for some reason. Here's the output log.

export: data=data/coco128.yaml, weights=['yolov5m_coco.pt'], imgsz=[640, 640], batch_size=256, device=0, half=False, inplace=False, keras=False, optimize=False, int8=False, dynamic=True, simplify=False, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, include=['engine']
fatal: detected dubious ownership in repository at '/media/dev/aditya/yolov5'
To add an exception for this directory, call:

	git config --global --add safe.directory /media/dev/aditya/yolov5
YOLOv5 🚀 2022-9-29 Python-3.7.13 torch-1.12.1+cu102 CUDA:0 (Quadro RTX 8000, 48601MiB)

Fusing layers... 
YOLOv5m summary: 290 layers, 21172173 parameters, 0 gradients

PyTorch: starting from yolov5m_coco.pt with output shape (256, 25200, 85) (40.8 MB)

False starting export with onnx 1.12.0...
[W shape_type_inference.cpp:425] Warning: Constant folding in symbolic shape inference fails: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select) (function ComputeConstantFolding)
[W shape_type_inference.cpp:425] Warning: Constant folding in symbolic shape inference fails: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select) (function ComputeConstantFolding)
[W shape_type_inference.cpp:425] Warning: Constant folding in symbolic shape inference fails: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select) (function ComputeConstantFolding)
ONNX: export failure ❌ 4.1s: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select)

TensorRT: starting export with TensorRT 8.4.3.1...
[10/04/2022-15:08:23] [TRT] [I] [MemUsageChange] Init CUDA: CPU +285, GPU +0, now: CPU 2321, GPU 20679 (MiB)
[10/04/2022-15:08:24] [TRT] [I] [MemUsageChange] Init builder kernel library: CPU +206, GPU +70, now: CPU 2544, GPU 20749 (MiB)
export.py:270: DeprecationWarning: Use set_memory_pool_limit instead.
  config.max_workspace_size = workspace * 1 << 30
[10/04/2022-15:08:24] [TRT] [I] ----------------------------------------------------------------
[10/04/2022-15:08:24] [TRT] [I] Input filename:   yolov5m_coco.onnx
[10/04/2022-15:08:24] [TRT] [I] ONNX IR version:  0.0.7
[10/04/2022-15:08:24] [TRT] [I] Opset version:    12
[10/04/2022-15:08:24] [TRT] [I] Producer name:    pytorch
[10/04/2022-15:08:24] [TRT] [I] Producer version: 1.12.1
[10/04/2022-15:08:24] [TRT] [I] Domain:           
[10/04/2022-15:08:24] [TRT] [I] Model version:    0
[10/04/2022-15:08:24] [TRT] [I] Doc string:       
[10/04/2022-15:08:24] [TRT] [I] ----------------------------------------------------------------
[10/04/2022-15:08:24] [TRT] [W] onnx2trt_utils.cpp:369: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
TensorRT: input "images" with shape(400, 3, 640, 640) DataType.FLOAT
TensorRT: output "output0" with shape(400, 25200, 85) DataType.FLOAT
TensorRT: building FP32 engine as yolov5m_coco.engine
export.py:297: DeprecationWarning: Use build_serialized_network instead.
  with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
[10/04/2022-15:08:24] [TRT] [E] 4: [network.cpp::operator()::3020] Error Code 4: Internal Error (images: kMIN dimensions in profile 0 are [1,3,640,640] but input has static dimensions [400,3,640,640].)
TensorRT: export failure ❌ 5.3s: __enter__

### Environment

_No response_

### Minimal Reproducible Example

_No response_

### Additional

_No response_

### Are you willing to submit a PR?

- [ ] Yes I'd like to help by submitting a PR!


Discussion:

Resolving patch:
diff --git a/export.py b/export.py
--- a/export.py
+++ b/export.py
@@ -251,11 +251,11 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
     if trt.__version__[0] == '7':  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
         grid = model.model[-1].anchor_grid
         model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
-        export_onnx(model, im, file, 12, False, dynamic, simplify)  # opset 12
+        export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
         model.model[-1].anchor_grid = grid
     else:  # TensorRT >= 8
         check_version(trt.__version__, '8.0.0', hard=True)  # require tensorrt>=8.0.0
-        export_onnx(model, im, file, 12, False, dynamic, simplify)  # opset 12
+        export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
     onnx = file.with_suffix('.onnx')
 
     LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
@@ -285,7 +285,7 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
 
     if dynamic:
         if im.shape[0] <= 1:
-            LOGGER.warning(f"{prefix}WARNING ⚠️ --dynamic model requires maximum --batch-size argument")
+            LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument")
         profile = builder.create_optimization_profile()
         for inp in inputs:
             profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)


Test patch: