Export.py need train mode export for cormel model to add NMS and decode layer

[junmcenroe]

Search before asking

• I have searched the YOLOv5 issues and found no similar bug report.

YOLOv5 Component

Export

Bug

I had had been converted with export.py with -train mode to convert cormel model. after exported .mlmodel, I added decode layer and NMS. But after eliminate -train mode of export.py, I found the following error.

---

/usr/local/lib/python3.7/dist-packages/coremltools/models/neural_network/builder.py in set_output(self, output_names, output_dims)
529 spec = self.spec
530 for idx, dim in enumerate(output_dims):
--> 531 spec.description.output[idx].type.multiArrayType.ClearField("shape")
532 spec.description.output[idx].type.multiArrayType.shape.extend(dim)
533 spec.description.output[

IndexError: list index (1) out of range

---

Exported model shape with -train model and eval mode are different. So I modify the export.py with train mode, no failure found.
On "Remove .train() mode exports #9429", this train() mode removed. but Could you add train() mode to export.py ?

Environment

Google Colab

Minimal Reproducible Example

No response

Additional

No response

Are you willing to submit a PR?

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

[github-actions]

👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.

Access additional YOLOv5 🚀 resources:

• Wiki – https://github.com/ultralytics/yolov5/wiki
• Tutorials – https://docs.ultralytics.com/yolov5
• Docs – https://docs.ultralytics.com

Access additional Ultralytics ⚡ resources:

• Ultralytics HUB – https://ultralytics.com/hub
• Vision API – https://ultralytics.com/yolov5
• About Us – https://ultralytics.com/about
• Join Our Team – https://ultralytics.com/work
• Contact Us – https://ultralytics.com/contact

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 YOLOv5 🚀 and Vision AI ⭐!

[glenn-jocher]

@junmcenroe hi there! Thanks for reaching out. It's great to see your interest in improving the export process. 😊

To clarify, the removal of .train() mode from export.py was intentional, as it wasn't used during inference and was causing inconsistencies in the exported models.

To address your requirement for the decode layer and NMS, you can customize your model by modifying the YOLO layer directly in the model.yaml configuration file before exporting. By doing this, you can include the desired layers without the need for the -train mode.

For more information, you can refer to our documentation at https://docs.ultralytics.com/yolov5/ for additional guidance on customizing the export process.

I appreciate your willingness to submit a PR! If you have any further questions or need assistance while working on the PR, feel free to reach out. We're here to help!

[junmcenroe]

Hi glenn-josher

Thanks for your supplement comment. As of today, your original export.py can export cormel formats with NMS. So no issue for me.
You and your teams' great continuous work save us.

[glenn-jocher]

@junmcenroe thank you for your kind words! The YOLO community and the Ultralytics team have indeed put in a lot of effort to enhance the YOLOv5 repository and we're glad to hear that our work has been helpful to you. If you have any further questions or need support in the future, please don't hesitate to reach out. Your feedback and contributions are highly valued. Keep up the great work, and happy coding!