Fix batchnorm inference, also for batch_size > 1

Author: Talmaj

README.md

@@ -34,9 +34,9 @@ Currently supported and tested models from [onnx_zoo](https://github.com/onnx/mo
 
 ## Limitations
 Known current version limitations are:
-- `batch_size > 1` could deliver unexpected results due to ambiguity of onnx's BatchNorm layer.   
-That is why in this case for now we raise an assertion error.  
-Set `experimental=True` in `ConvertModel` to be able to use `batch_size > 1`.
+- `batch_size > 1` requires `experimental=True` in `ConvertModel`.   
+BatchNorm layers use inference mode (running statistics), which is correct for most ONNX models  
+exported for inference. If your model was exported in training mode, results may differ.
 - Fine tuning and training of converted models was not tested yet, only inference.
 
 ## Development

onnx2pytorch/convert/model.py

@@ -96,8 +96,9 @@ def __init__(
         batch_dim: int
             Dimension of the batch.
         experimental: bool
-            Experimental implementation allows batch_size > 1. However,
-            batchnorm layers could potentially produce false outputs.
+            Experimental implementation allows batch_size > 1.
+            BatchNorm layers use inference mode (running statistics), which is
+            correct for most ONNX models exported for inference.
         enable_pruning: bool
             Track kept/pruned indices between different calls to forward pass.
 
@@ -145,8 +146,11 @@ def __init__(
 
         if experimental:
             warnings.warn(
-                "Using experimental implementation that allows 'batch_size > 1'."
-                "Batchnorm layers could potentially produce false outputs."
+                "Using experimental implementation that allows 'batch_size > 1'. "
+                "BatchNorm layers use inference mode (running statistics). "
+                "This is correct for ONNX models exported for inference, but may be "
+                "incorrect if the model was exported in training mode.",
+                UserWarning,
             )
 
     def forward(self, *input_list, **input_dict):

onnx2pytorch/operations/batchnorm.py

@@ -48,6 +48,9 @@ def __init__(self, torch_params, *args, **kwargs):
             for key, value in zip(keys, torch_params):
                 getattr(self.bnu, key).data = value
 
+        # Set to eval mode to use running statistics (ONNX inference behavior)
+        self.bnu.eval()
+
     def forward(self, X, scale=None, B=None, input_mean=None, input_var=None):
         if self.has_lazy:
             self.bnu.initialize_parameters(X)