Fix ONNX export using --grid --simplify --dynamic simultaneously

[jylink]

python models/export.py --grid --dynamic --simplify failed to export onnx model.

#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

_{Made with ❤️ by Ultralytics Actions}

🌟 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. 🌐

[github-actions]

👋 Hello @jylink, thank you for submitting a 🚀 PR! To allow your work to be integrated as seamlessly as possible, we advise you to:

• ✅ Verify your PR is up-to-date with origin/master. If your PR is behind origin/master an automatic GitHub actions rebase may be attempted by including the /rebase command in a comment body, or by running the following code, replacing 'feature' with the name of your local branch:

git remote add upstream https://github.com/ultralytics/yolov5.git
git fetch upstream
git checkout feature  # <----- replace 'feature' with local branch name
git rebase upstream/master
git push -u origin -f

• ✅ Verify all Continuous Integration (CI) checks are passing.
• ✅ Reduce changes to the absolute minimum required for your bug fix or feature addition. "It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is." -Bruce Lee

[glenn-jocher]

@jylink thanks for the PR! I've just now merged a similar modification to the Detect() layer in 'YOLOv5 AWS Inferentia Inplace compatibility updates' #2953 which implements some of the same updates and will cause a conflict with this PR.

Once #2953 is merged please merge master in your branch, resolving conflicts as they appear. You may also be able to recycle the new Detect layer inplace attribute.

[glenn-jocher]

@jylink this looks a lot better, thanks for the updates! Will review later today.

[glenn-jocher]

@jylink I've taken the opportunity here to streamline the export a bit:

• I've removed opt.grid completely, new default is to always output grid as I don't see many use-cases without it.
• I've added your Detect() layer attribute assignments into the if statement directly above them.
• I've added a new opt.inplace argument which simply updates the existing Detect.inplace value. Default is inplace=False, passing --inplace sets it true.

Can you verify these 3 use cases on your side and let me know what you think?

python models/export.py --dynamic
python models/export.py --simplify
python models/export.py --dynamic --simplify

[glenn-jocher]

@jylink I think perhaps the -1 in the .view() call is causing problems in yolo.py L38. Can you try replacing these lines with the explicit shape below and see 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. If this solution works you can revert your yolo.py L60 addition.

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)

[glenn-jocher]

@jylink also I think your L60 addition may have a bug, as anchor_grid may never assume a value other than 1 in dimension 1, whereas you have bs in dimension 1, which is the batch-size. So this will likely crash on non-unity batch-sizes I think.

[jylink]

@jylink I've taken the opportunity here to streamline the export a bit:

• I've removed opt.grid completely, new default is to always output grid as I don't see many use-cases without it.
• I've added your Detect() layer attribute assignments into the if statement directly above them.
• I've added a new opt.inplace argument which simply updates the existing Detect.inplace value. Default is inplace=False, passing --inplace sets it true.

Can you verify these 3 use cases on your side and let me know what you think?

python models/export.py --dynamic
python models/export.py --simplify
python models/export.py --dynamic --simplify

They are ok

@jylink I think perhaps the -1 in the .view() call is causing problems in yolo.py L38. Can you try replacing these lines with the explicit shape below and see 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. If this solution works you can revert your yolo.py L60 addition.

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)

Simplify failed
ONNX: simplifier failure: [ONNXRuntimeError] : 1 : FAIL : Node (Mul_382) Op (Mul) [ShapeInferenceError] Incompatible dimensions
Even without simplify, onnxruntime failed as well (since grid is default)
Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from ../best.onnx failed:Node (Mul_382) Op (Mul) [ShapeInferenceError] Incompatible dimensions

@jylink also I think your L60 addition may have a bug, as anchor_grid may never assume a value other than 1 in dimension 1, whereas you have bs in dimension 1, which is the batch-size. So this will likely crash on non-unity batch-sizes I think.

Changed bs to 1, passed all

[glenn-jocher]

@jluntamazon 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 an explicit view shape which it seems to lack otherwise on export.
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
to:
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2) # wh

[jluntamazon]

@jluntamazon would this PR update to this line affect your use-case?

That should be fine for our use-case, the main operation I was avoiding in the original PR was slice assignment

[glenn-jocher]

@jluntamazon thanks for the feedback!

@jylink PR is merged! Thank you for your contributions.