[Relax] Batch norm correctness on eval mode #17752
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
hugolatendresse
changed the title
Contributor
Author
|
cc: @MasterJH5574 this is ready for review |
Contributor
Author
|
@tvm-bot rerun |
MasterJH5574
approved these changes
Mar 26, 2025
Contributor
MasterJH5574
left a comment
MasterJH5574
left a comment
There was a problem hiding this comment.
Looks good. Thank you @hugolatendresse for the enhancement!
| ########## Neural Network ########## | ||
|
|
||
| def _batch_norm_legit_no_training(self, node: fx.Node) -> relax.Var: | ||
| def _batch_norm(self, node: fx.Node, training) -> relax.Var: |
Contributor
There was a problem hiding this comment.
Good to add a type annotation in any of followup PRs.
Suggested change
| def _batch_norm(self, node: fx.Node, training) -> relax.Var: | |
| def _batch_norm(self, node: fx.Node, training: bool) -> relax.Var: |
Contributor
Author
There was a problem hiding this comment.
Got it, will do, thanks
ShiboXing
pushed a commit
to ShiboXing/tvm
that referenced
this pull request
Aug 10, 2025
Batch_norm is a different operator in training and eval. The previous interface defaulted to the training mode and required changing an ingested pytorch program itself to use the eval mode. This is sub-ideal, especially since torch.export explicitely communicates whether batch_norm should be in training or eval in a given torch program. This PR automates the selection of training/eval mode in the exported program translator, and achieves correctness for eval mode. Future TODO: there is something wrong with batch_norm on training mode. It does not pass a correctness test when taken straight from the main branch (there's an issue with tensor dimensions). I added a note to address later as training mode is probably not high priority.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Batch_norm is a different operator in training and eval. The previous interface defaulted to the training mode and required changing an ingested pytorch program itself to use the eval mode. This is sub-ideal, especially since torch.export explicitely communicates whether batch_norm should be in training or eval in a given torch program.
This PR automates the selection of training/eval mode in the exported program translator, and achieves correctness for eval mode.
Future TODO: there is something wrong with batch_norm on training mode. It does not pass a correctness test when taken straight from the main branch (there's an issue with tensor dimensions). I added a note to address later as training mode is probably not high priority.