Repository: keras-team/keras

Issue:
BatchNormalization layer fails with torch backend on GPU with `scale=False, center=True`
Here is the colab for reproducing it.
[link](https://colab.research.google.com/github/haifeng-jin/Colabs/blob/main/Keras_torch_BN.ipynb)


Discussion:
The bug is fixed by https://github.com/keras-team/keras/commit/9815ac1c81f0f24cf3ca24638e4ab11c5fc95bd1.
However, the torch BN is converging slower than other backends. Needs further investigation.

Resolving patch:
diff --git a/keras/backend/torch/nn.py b/keras/backend/torch/nn.py
--- a/keras/backend/torch/nn.py
+++ b/keras/backend/torch/nn.py
@@ -684,7 +684,7 @@ def moments(x, axes, keepdims=False, synchronized=False):
     # gradient is zero.
     variance = torch.mean(
         torch.square(x), dim=axes, keepdim=True
-    ) - torch.square(mean.detach())
+    ) - torch.square(mean)
 
     if not keepdims:
         mean = torch.squeeze(mean, axes)
@@ -710,45 +710,30 @@ def batch_normalization(
     x, mean, variance, axis, offset=None, scale=None, epsilon=1e-3
 ):
     x = convert_to_tensor(x)
-    mean = convert_to_tensor(mean).detach()
-    variance = convert_to_tensor(variance).detach()
+    mean = convert_to_tensor(mean)
+    variance = convert_to_tensor(variance)
+
+    shape = [1] * len(x.shape)
+    shape[axis] = mean.shape[0]
+    mean = torch.reshape(mean, shape)
+    variance = torch.reshape(variance, shape)
+
     if offset is not None:
         offset = convert_to_tensor(offset)
+        offset = torch.reshape(offset, shape)
     else:
         offset = torch.zeros_like(mean)
     if scale is not None:
         scale = convert_to_tensor(scale)
+        scale = torch.reshape(scale, shape)
     else:
         scale = torch.ones_like(variance)
 
-    def _batch_norm():
-        return tnn.batch_norm(
-            input=x,
-            running_mean=mean,
-            running_var=variance,
-            weight=scale,
-            bias=offset,
-            training=False,
-            eps=epsilon,
-        )
-
-    if axis == 1:
-        return _batch_norm()
-
-    if axis < 0:
-        axis = len(x.shape) + axis
-
-    order = list(range(len(x.shape)))
-    order.pop(axis)
-    order.insert(1, axis)
-    x = x.permute(order)
-
-    x = _batch_norm()
-
-    order = list(range(len(x.shape)))
-    order.pop(1)
-    order.insert(axis, 1)
-    return x.permute(order)
+    return (
+        x.subtract(mean)
+        .mul_(variance.add(epsilon).rsqrt_().mul(scale))
+        .add_(offset)
+    )
 
 
 def ctc_loss(


Test patch:
diff --git a/integration_tests/numerical_test.py b/integration_tests/numerical_test.py
--- a/integration_tests/numerical_test.py
+++ b/integration_tests/numerical_test.py
@@ -36,10 +36,12 @@ def build_keras_model(keras_module, num_classes):
             keras_module.layers.Conv2D(
                 32, kernel_size=(3, 3), activation="relu"
             ),
+            keras_module.layers.BatchNormalization(),
             keras_module.layers.MaxPooling2D(pool_size=(2, 2)),
             keras_module.layers.Conv2D(
                 64, kernel_size=(3, 3), activation="relu"
             ),
+            keras_module.layers.BatchNormalization(scale=False, center=True),
             keras_module.layers.MaxPooling2D(pool_size=(2, 2)),
             keras_module.layers.Flatten(),
             keras_module.layers.Dense(num_classes, activation="softmax"),