Training mode removal from ONNX nodes

[i-riyad]

What does this PR do?

Type of change: Bug fix

Overview: Training mode removal from a ONNX node was incomplete.

Usage

from modelopt.onnx.utils import remove_node_training_mode

result_model = remove_node_training_mode(model, "BatchNormalization")

Testing

python -m pytest tests/unit/torch/deploy/utils/test_torch_onnx_utils.py::test_remove_node_training_mode_attribute -v
python -m pytest tests/unit/torch/deploy/utils/test_torch_onnx_utils.py::test_remove_node_extra_training_outputs -v

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes
• Did you write any new necessary tests?: Yes
• Did you add or update any necessary documentation?: No
• Did you update Changelog?: No

Additional Information

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[ajrasane]

Is there a way we can check that that this output that we remove will be the training output?

[i-riyad]

Actually, the number of training outputs is operator specific and removing them specifically is a bit more convoluted. I have implemented the same functionality by removing all the unused outputs from a given node if training_mode is on. Please review.

[ajrasane]

Could you add a comment to show how this model will look after creation?

[ajrasane]

Could you also add another test to check if there are multiple outputs of a node that are not training outputs, they are not removed?

[i-riyad]

Modified the existing test to cover this.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.97%. Comparing base (dba0b37) to head (ff5f6df).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #277      +/-   ##
==========================================
+ Coverage   73.83%   73.97%   +0.14%     
==========================================
  Files         173      172       -1     
  Lines       17402    17351      -51     
==========================================
- Hits        12849    12836      -13     
+ Misses       4553     4515      -38     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.