Export torchscript with NMS

[RobinFrcd]

Search before asking

• 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! 😊