Handle multiple dimension case for BatchNormalization (#106)

* Handle multiple dimension case for BatchNormalization

* Handle multiple dimension case for BatchNormalization

Author: jiafatom

keras2onnx/ke2onnx/batch_norm.py

@@ -10,11 +10,18 @@
 
 def convert_keras_batch_normalization(scope, operator, container):
     op = operator.raw_operator
-    if (op.axis != 3 and op.axis != -1) or len(op.input_shape) == 2:
+    skip_transpose = (op.axis != len(op.input_shape) - 1 and op.axis != -1) or len(op.input_shape) <= 2
+    if not skip_transpose:
+        perm_1 = list(range(1, len(op.input_shape) - 1))
+        perm_1 = [0, len(op.input_shape) - 1] + perm_1
+        perm_2 = list(range(2, len(op.input_shape)))
+        perm_2 = [0] + perm_2 + [1]
+
+    if skip_transpose:
         adjusted_input_name = operator.inputs[0].full_name
     else:
         adjusted_input_name = scope.get_unique_variable_name(operator.inputs[0].full_name + '_transposed')
-        apply_transpose(scope, operator.inputs[0].full_name, adjusted_input_name, container, perm=[0, 3, 1, 2])
+        apply_transpose(scope, operator.inputs[0].full_name, adjusted_input_name, container, perm=perm_1)
 
     input_tensor_names = [adjusted_input_name]
 
@@ -49,7 +56,7 @@ def convert_keras_batch_normalization(scope, operator, container):
     momentum = op.momentum
     spatial = 1
 
-    if (op.axis != 3 and op.axis != -1) or len(op.input_shape) == 2:
+    if skip_transpose:
         # If no transpose is required, we can simply use the output of ONNX BatchNorm as the final outcome
         apply_batch_norm(scope, input_tensor_names, operator.output_full_names, container,
                          operator_name=operator.full_name, epsilon=epsilon, is_test=is_test,
@@ -61,5 +68,5 @@ def convert_keras_batch_normalization(scope, operator, container):
                          operator_name=operator.full_name, epsilon=epsilon, is_test=is_test,
                          momentum=momentum, spatial=spatial)
 
-        # Permute [N,C,H,W] to [N,H,W,C]
-        apply_transpose(scope, intermediate_output_name, operator.outputs[0].full_name, container, perm=[0, 2, 3, 1])
+        # For 4D case, this is to permute [N,C,H,W] to [N,H,W,C]
+        apply_transpose(scope, intermediate_output_name, operator.outputs[0].full_name, container, perm=perm_2)

tests/test_layers.py

@@ -552,17 +552,38 @@ def test_batch_normalization(self):
         self._batch_norm_helper(data, 'zeros', 'zeros', False, True, 1)
 
     def test_batch_normalization_2(self):
-        # test batch normalization on 2D input
-        input_dim = 10
-        batch_size = 4
-        model = keras.models.Sequential()
-        model.add(keras.layers.InputLayer(input_shape=(input_dim,)))
-        model.add(keras.layers.BatchNormalization(axis=-1))
-        model.add(keras.layers.Dense(5))
-        data = np.random.randn(batch_size, input_dim).astype(np.float32)
-        onnx_model = keras2onnx.convert_keras(model)
-        expected = model.predict(data)
-        self.assertTrue(self.run_onnx_runtime('test_batch_normalization_2', onnx_model, [data], expected))
+        for axis in [1, -1]:
+            batch_size = 4
+            input_dim_1 = 10
+            input_dim_2 = 20
+            input_dim_3 = 30
+
+            model = keras.models.Sequential()
+            model.add(keras.layers.InputLayer(input_shape=(input_dim_1,)))
+            model.add(keras.layers.BatchNormalization(axis=axis))
+            model.add(keras.layers.Dense(5))
+            data = np.random.randn(batch_size, input_dim_1).astype(np.float32)
+            onnx_model = keras2onnx.convert_keras(model)
+            expected = model.predict(data)
+            self.assertTrue(self.run_onnx_runtime('test_batch_normalization_2_2d', onnx_model, [data], expected))
+
+            model = keras.models.Sequential()
+            model.add(keras.layers.InputLayer(input_shape=(input_dim_1, input_dim_2)))
+            model.add(keras.layers.BatchNormalization(axis=axis))
+            model.add(keras.layers.Dense(5))
+            data = np.random.randn(batch_size, input_dim_1, input_dim_2).astype(np.float32)
+            onnx_model = keras2onnx.convert_keras(model)
+            expected = model.predict(data)
+            self.assertTrue(self.run_onnx_runtime('test_batch_normalization_2_3d', onnx_model, [data], expected))
+
+            model = keras.models.Sequential()
+            model.add(keras.layers.InputLayer(input_shape=(input_dim_1, input_dim_2, input_dim_3)))
+            model.add(keras.layers.BatchNormalization(axis=axis))
+            model.add(keras.layers.Dense(5))
+            data = np.random.randn(batch_size, input_dim_1, input_dim_2, input_dim_3).astype(np.float32)
+            onnx_model = keras2onnx.convert_keras(model)
+            expected = model.predict(data)
+            self.assertTrue(self.run_onnx_runtime('test_batch_normalization_2_4d', onnx_model, [data], expected))
 
     def test_simpleRNN(self):
         inputs1 = keras.Input(shape=(3, 1))