Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update export.py with --train mode argument #3066

Merged
merged 1 commit into from
May 7, 2021

Conversation

glenn-jocher
Copy link
Member

@glenn-jocher glenn-jocher commented May 7, 2021

This will allow error-free export of YOLOv5 to CoreML models. All ops are included except for the Detect layer grid reconstruction, which is not used during training. The model will output the 3 or 4 output layers rather than the concatenated reconstructed grid output used during .eval() mode.

🛠️ PR Summary

Made with ❤️ by Ultralytics Actions

🌟 Summary

Enhanced export script with training mode option 🛠️

📊 Key Changes

  • Added an argument (--train) to the export script allowing users to specify training mode during export.
  • Ensured compatibility with PyTorch 1.6.0 by setting non-persistent buffer sets.

🎯 Purpose & Impact

  • Purpose: The addition of --train allows the model to be exported in training mode, which is useful for continuing training or fine-tuning on other datasets.
  • Impact: Users will have more flexibility in how they export models, potentially improving the model's adaptability to various tasks. It also assures better forward compatibility with older versions of PyTorch.

@glenn-jocher
Copy link
Member Author

Should resolve #2961

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant