ONNX export in .train() mode fix

[ChaofWang]

This PR is submitted for this bug #3346

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Enhanced ONNX export options for different training modes.

📊 Key Changes

• 🔄 Added an option to specify training or evaluation mode during ONNX export.
• 🧱 Enabled or disabled constant folding depending on the selected mode (training or evaluation).

🎯 Purpose & Impact

• 🎛 This improvement allows developers to export their models in the proper mode, matching their intended use case (either for further training or for inference).
• 🚀 Users can now optimize their models' ONNX export for training with gradient updates or for faster inference without gradients.
• 🛠 The option to toggle constant folding can help in optimizing the model accordingly, potentially improving performance in inference scenarios.

[github-actions]

👋 Hello @ChaofWang, thank you for submitting a 🚀 PR! To allow your work to be integrated as seamlessly as possible, we advise you to:

• ✅ Verify your PR is up-to-date with origin/master. If your PR is behind origin/master an automatic GitHub actions rebase may be attempted by including the /rebase command in a comment body, or by running the following code, replacing 'feature' with the name of your local branch:

git remote add upstream https://github.com/ultralytics/yolov5.git
git fetch upstream
git checkout feature  # <----- replace 'feature' with local branch name
git rebase upstream/master
git push -u origin -f

• ✅ Verify all Continuous Integration (CI) checks are passing.
• ✅ Reduce changes to the absolute minimum required for your bug fix or feature addition. "It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is." -Bruce Lee

[glenn-jocher]

@ChaofWang PR is merged, thank you for your contributions!