[ONNX] Add training mode support for BatchNormalization op

[vivekkhandelwal1]

This commit extends the OnnxToTorch lowering for BatchNormalization op for supporting the case when training=True.

Signed-Off By: Vivek Khandelwal vivekkhandelwal1424@gmail.com

[vivekkhandelwal1]

Fixes nod-ai/SHARK-ModelDev#285 (comment)

[zjgarvey]

This does not pass e2e testing for the following onnx node tests due to a significant numerics mismatch:

"test_batchnorm_epsilon_training_mode"
"test_batchnorm_example_training_mode"

Please try to debug the numerics mismatching when you get the chance. You can either use this branch of torch-mlir in your iree-build and run the iree_tests, or you can use the alt_e2eshark and run python run.py --torchtolinalg -t test_batch.

[zjgarvey]

From a glance, it looks like the numerics for Y and running_var are off, but running_mean seems to match pretty well.

[vivekkhandelwal1]

This does not pass e2e testing for the following onnx node tests due to a significant numerics mismatch:

"test_batchnorm_epsilon_training_mode" "test_batchnorm_example_training_mode"

Please try to debug the numerics mismatching when you get the chance. You can either use this branch of torch-mlir in your iree-build and run the iree_tests, or you can use the alt_e2eshark and run python run.py --torchtolinalg -t test_batch.

Hi @zjgarvey, the accuracy issue is fixed. After a lot of debugging, it turned out that the unbiased has to be set to False, instead of True.

[vivekkhandelwal1]

@zjgarvey Can you please review this PR, today?

[zjgarvey]

As long as this is passing numerics, the implementation is clear and well-commented on. Thanks, Vivek.