search

ultralytics / yolov5

YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite
https://docs.ultralytics.com
GNU Affero General Public License v3.0
50.87k stars 16.38k forks source link

Export torchscript with NMS #13308

Open RobinFrcd opened 2 months ago

RobinFrcd commented 2 months ago

Search before asking

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?

UltralyticsAssistant commented 2 months ago

🤖 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:

Thanks for being a part of the Ultralytics community! 😊

pderrenger commented 5 days ago

@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.