Repository: ultralytics/yolov5
Pull request: Fix ONNX export using --grid --simplify --dynamic simultaneously
`python models/export.py --grid --dynamic --simplify` failed to export onnx model.

[#2856](https://github.com/ultralytics/yolov5/pull/2856) fix dynamic and simplify but still cannot work with grid, because the `self.grid[i]` mismatch the dynamic `y[..., 0:2]`. This pull request fixes it.

Also the removal of subscript assignments helps to avoid unsupported ScatterND nodes, so the onnx exported using `--grid` can be directly converted to tensorRT engine now

## 🛠️ PR Summary

<sub>Made with ❤️ by [Ultralytics Actions](https://github.com/ultralytics/actions)<sub>

### 🌟 Summary
Improved ONNX export flexibility for YOLOv5.

### 📊 Key Changes
- Removed the `--grid` flag from the command-line arguments in `export.py`.
- Added a new `--inplace` flag to control the `inplace` parameter for the YOLOv5 Detect() layer.
- Modified how the Detect() layer handles the export grid and the use of inplace operations.
- Updated the logic for dynamic axis computation in ONNX exports and grid generation during inference.

### 🎯 Purpose & Impact
- The changes aim to optimize the export process of YOLOv5 models for different use cases, such as dynamic batching.
- The `--inplace` flag provides users with the option to choose how Detect() processes inputs, which could lead to memory optimizations. 
- These adjustments may improve the model's compatibility with varying deployment environments, potentially enhancing performance and facilitating integration with platforms like AWS Inferentia. 🚀
- Users leveraging ONNX models should enjoy more flexibility and potentially better inference speed and device compatibility. 🌐

Problem statement:
Export with ONNX Simplifier with --grid error
## 🐛 Bug
An exported model as ONNX using --grid parameter cannot be used by onnx-runtime or simplified by onnx-simplifier
A Mul Node triggers a shape inference error Incompatible dimensions


## To Reproduce

Replace ONNX export in export.py with this code and run with command `python3 models/export.py --grid`
```py
try:
        import onnx
        from onnxsim import simplify
        print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
        f = opt.weights.replace('.pt', '.onnx')  # filename
        torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
                          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)

        # Checks
        onnx_model = onnx.load(f)  # load onnx model
        onnx.checker.check_model(onnx_model)  # check onnx model

        # This step triggers the error
        model_simp, check = simplify(onnx_model)
        onnx.save(model_simp, f)

        # print(onnx.helper.printable_graph(onnx_model.graph))  # print a human readable model
        print('ONNX export success, saved as %s' % f)
    except Exception as e:
        print('ONNX export failure: %s' % e)
```

Output:
```
Starting ONNX export with onnx 1.8.1...
ONNX export failure: [ONNXRuntimeError] : 1 : FAIL : Node (Mul_925) Op (Mul) [ShapeInferenceError] Incompatible dimensions
```


## Expected behavior
Any yolov5 model exported as ONNX should be valid

Hints from issue discussion:
@antlamon thanks for the bug report. We don't generally provide support for code customizations and external package not in requirements.txt. 

If an external package is causing an error you may also want to raise an issue with the package authors.
I would like to add that without any modifications to export.py, the --grid option also results in an unusable .onnx file when ran on the yolov5s.pt model. It works fine without the --grid option 

During the running of the script, the following warning was produced (By torchscipt, not onnx though):

```
Starting TorchScript export with torch 1.8.0+cu101...
./models/yolo.py:48: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if self.grid[i].shape[2:4] != x[i].shape[2:4]:
/usr/local/lib/python3.7/dist-packages/torch/jit/_trace.py:940: TracerWarning: Encountering a list at the output of the tracer might cause the trace to be incorrect, this is only valid if the container structure does not change based on the module's inputs. Consider using a constant container instead (e.g. for `list`, use a `tuple` instead. for `dict`, use a `NamedTuple` instead). If you absolutely need this and know the side effects, pass strict=False to trace() to allow this behavior.
  _force_outplace,
TorchScript export success, saved as yolov5s.torchscript.pt
```
Attempting to run a inference session results in 

```
Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from yolov5s.onnx failed:Node (Mul_925) Op (Mul) [ShapeInferenceError] Incompatible dimensions```
This error occurred to me when exporting the onnx model using torch==1.8.1 with torchvision==0.9.1. When i export using torch==1.7.1, the loading of the onnx model works fine in both torch==1.7.1 and torch==1.8.1.
Thank you for pointing that out. That indeed was the issue
Will there be a fix, since torchvision==0.8.2 (required by torch 1.7.1) doesn't exist for windows?

I am also getting the same error,downgrading torch and torchvision versions didn't help me out to fix this issue.

When I downgrade the pytorch version and export with --dynamic --grid, I can load the model, but it fails when doing inference on a (1, 3, 1088, 1920) tensor with this:

2021-04-15 20:26:39.079097920 [E:onnxruntime:, sequential_executor.cc:339 Execute] Non-zero status code returned while running Add node. Name:'Add_945' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/math/element_wise_ops.h:487 void onnxruntime::BroadcastIterator::Append(ptrdiff_t, ptrdiff_t) axis == 1 || axis == largest was false. Attempting to broadcast an axis by a dimension other than 1. 34 by 60

Traceback (most recent call last):
  File "test_onnx.py", line 42, in <module>
    outputs = sess.run(None, {input_name: image})
  File "venv/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 188, in run
    return self._sess.run(output_names, input_feed, run_options)
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Add node. Name:'Add_945' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/math/element_wise_ops.h:487 void onnxruntime::BroadcastIterator::Append(ptrdiff_t, ptrdiff_t) axis == 1 || axis == largest was false. Attempting to broadcast an axis by a dimension other than 1. 34 by 60

It does work with just --dynamic. I've seen some approaches where people re-implement the last layer with onnx. I guess that's probably the best approach for now.
@antlamon @Lucashsmello @thestonehead @timstokman we've integrated onnx-simplifier into export.py now in ONNX Simplifier PR #2815 and verified it's passing CI on all operating systems.

I'm not sure if this resolves the original issue, but hopefully it's a step in the right direction.
> @antlamon @Lucashsmello @thestonehead @timstokman we've integrated onnx-simplifier into export.py now in ONNX Simplifier PR #2815 and verified it's passing CI on all operating systems.
> 
> I'm not sure if this resolves the original issue, but hopefully it's a step in the right direction.

@glenn-jocher Unfortunately, the problem still persists: I am using the docker image (version v5.0) and `--grid` causes ONNX export to fail on simplifying. The resulting onnx file cannot be used due to `Incompatible dimensions` error. However, rolling back pytorch to 1.8.0 (i.e. using the docker image v4.0 with latest repository version, which includes ONNX simplifier) works OK.

Commit message: rename exp_dynamic to onnx_dynamic, comment

Patch to review:
diff --git a/models/export.py b/models/export.py
index da15079149a1..44bf1c96d959 100644
--- a/models/export.py
+++ b/models/export.py
@@ -26,8 +26,8 @@
     parser.add_argument('--weights', type=str, default='./yolov5s.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('--grid', action='store_true', help='export Detect() layer grid')
     parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument('--inplace', action='store_true', help='inplace ops')  # inplace = True
     parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')  # ONNX-only
     parser.add_argument('--simplify', action='store_true', help='simplify ONNX model')  # ONNX-only
     opt = parser.parse_args()
@@ -56,9 +56,11 @@
                 m.act = Hardswish()
             elif isinstance(m.act, nn.SiLU):
                 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
+        elif isinstance(m, models.yolo.Detect):
+            m.inplace = opt.inplace
+            m.onnx_dynamic = opt.dynamic
+            # m.forward = m.forward_export  # assign forward (optional)
+
     for _ in range(2):
         y = model(img)  # dry runs
     print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)")
diff --git a/models/yolo.py b/models/yolo.py
index 520047ff3a99..4f999cedb5fc 100644
--- a/models/yolo.py
+++ b/models/yolo.py
@@ -24,7 +24,7 @@
 
 class Detect(nn.Module):
     stride = None  # strides computed during build
-    export = False  # onnx export
+    onnx_dynamic = False  # ONNX export parameter
 
     def __init__(self, nc=80, anchors=(), ch=(), inplace=True):  # detection layer
         super(Detect, self).__init__()
@@ -42,14 +42,13 @@ def __init__(self, nc=80, anchors=(), ch=(), inplace=True):  # detection layer
     def forward(self, x):
         # x = x.copy()  # for profiling
         z = []  # inference output
-        self.training |= self.export
         for i in range(self.nl):
             x[i] = self.m[i](x[i])  # conv
             bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
             x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
 
             if not self.training:  # inference
-                if self.grid[i].shape[2:4] != x[i].shape[2:4]:
+                if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
                     self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
 
                 y = x[i].sigmoid()
@@ -58,7 +57,7 @@ def forward(self, x):
                     y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                 else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
                     xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
-                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
+                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(bs, self.na, 1, 1, 2)  # wh
                     y = torch.cat((xy, wh, y[..., 4:]), -1)
                 z.append(y.view(bs, -1, self.no))
 


Code review comments:
diff_hunk: @@ -58,7 +57,7 @@ def forward(self, x):
                     y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                 else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
                     xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
-                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
+                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(bs, self.na, 1, 1, 2)  # wh
path: models/yolo.py
text: @user2:
@user1 would this PR update to this line affect your use-case? The change would be applied to yolo.py L61 and add a `.view()` call. The .view() operation does NOT change the shape of `wh`, the reason for this is that ONNX export requires this otherwise it sees ambiguity in the shape.
`                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh`
to:
`                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(bs, self.na, 1, 1, 2)  # wh`

@user2:
@author would it be possible to move this .view() call to the self.anchor_grid[i] definition in the init() function? Oh nevermind, I see it already has a .view() call. Maybe the -1 in the .view() call is causing problems. Can you try replacing these lines and seeing if this fixes the problem? Ideally we do not want to add operations in the forward method as these will run on every image, which is expensive.

Replace:
`self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2))  # shape(nl,1,na,1,1,2)`
with:
`self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, self.na, 1, 1, 2))  # shape(nl,1,na,1,1,2)`

@author:
@user3-jocher Hi, I've tried and it failed, you could find the error info in my previous reply

@user2:
Yes I see that, thanks! I'm leaving the PR as-is, just waiting on @user1 to confirm that this addition doesn't affect the recent AWS PR https://github.com/ultralytics/yolov5/pull/2953