Export torchscript with NMS

[RobinFrcd]

trafficstars

Search before asking

• [X] I have searched the YOLOv5 issues and found no similar feature requests.

Description

In the export script it's possible to export a coreml or tflite with nms but not for the default torchscript.

I want to export the full model (with pre-post processing functions) with torchscript. Is there any official recommended way to go or do we need to rely on slightly deprecated tools like yolort ?

Use case

Export YOLOv5 to torchscript with pre and post processing functions included in the torchscript.

Additional

No response

Are you willing to submit a PR?

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

[UltralyticsAssistant]

---

🤖 Automated Response

Hi @RobinFrcd, thanks for your interest in YOLOv5 and for taking the time to submit your issue!

It looks like you're looking to export YOLOv5 as a TorchScript model with pre and post-processing included. At the moment, the current export script in YOLOv5 does not natively support fully integrated pre-post processing for TorchScript, as it does for CoreML or TFLite.

Here's what you can do in the meantime:

1. Custom Post-Processing: You can modify the model's forward function to include post-processing steps manually. This can be complex, so make sure to thoroughly test your modifications.

2. Explore Alternatives: While not officially maintained by us, yolort might provide some insights or tools that could assist, but be mindful of its maintenance status.

3. Community Solutions: Check the discussions and issues for similar problems others might have solved.

🛠️ Contributions Welcome: If you're interested in experimenting and think you can improve the export process, consider submitting a Pull Request! We welcome contributions that align with best practices and project goals.

For more information on best practices for contribution, see our CONTRIBUTING guide.

A human reviewer will take a look at this issue as soon as possible. Your patience is appreciated!

🔗 Links:

• Export Script Reference
• YOLOv5 Documentation

Thanks for being a part of the Ultralytics community! 😊

[pderrenger]

@RobinFrcd currently, exporting YOLOv5 to TorchScript with integrated NMS is not officially supported. You may need to implement custom modifications or explore third-party solutions like yolort, though these might not be fully up-to-date. Keep an eye on the repository for any updates or new features that might address this in the future.

[github-actions[bot]]

👋 Hello there! We wanted to give you a friendly reminder that this issue has not had any recent activity and may be closed soon, but don't worry - you can always reopen it if needed. If you still have any questions or concerns, please feel free to let us know how we can help.

For additional resources and information, please see the links below:

• Docs: https://docs.ultralytics.com
• HUB: https://hub.ultralytics.com
• Community: https://community.ultralytics.com

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLO 🚀 and Vision AI ⭐