Added tests for FRN on N-D tensors, fixes tensorflow#1441

* Tests for 3D-7D tensors
* Error raised when tensors aren't at least 3D

Author: Opletts

tensorflow_addons/layers/normalizations.py

@@ -354,7 +354,7 @@ class FilterResponseNormalization(tf.keras.layers.Layer):
         (tuple of integers, does not include the samples axis)
         when using this layer as the first layer in a model. This layer supports
         arbitrary tensors with the following assumptions:
-            -   Expected input tensor to be at least 2D.
+            -   Expected input tensor to be at least 3D.
             -   0th index in tensor shape is expected to be the batch dimension.
 
     Output shape
@@ -384,7 +384,7 @@ def __init__(
     ):
         super().__init__(name=name, **kwargs)
         self.channel_idx = channel_idx
-        self.epsilon = tf.math.abs(tf.cast(epsilon, dtype=self.dtype))
+        self.epsilon = epsilon
         self.beta_initializer = tf.keras.initializers.get(beta_initializer)
         self.gamma_initializer = tf.keras.initializers.get(gamma_initializer)
         self.beta_regularizer = tf.keras.regularizers.get(beta_regularizer)
@@ -421,7 +421,7 @@ def build(self, input_shape):
         super().build(input_shape)
 
     def call(self, inputs):
-        epsilon = self.epsilon
+        epsilon = tf.math.abs(tf.cast(self.epsilon, dtype=self.dtype))
         if self.use_eps_learned:
             epsilon += tf.math.abs(self.eps_learned)
         nu2 = tf.reduce_mean(tf.square(inputs), axis=self.axis, keepdims=True)
@@ -484,11 +484,6 @@ def _check_axis(self, axis):
             self.axis = axis
 
         elif isinstance(axis, int):
-            if abs(axis) != 1 or abs(self.channel_idx) != 1:
-                raise ValueError(
-                    "Expected index for 2D is -1/1 but got {}".format(axis)
-                )
-
             self.axis = [axis]
 
         else:
@@ -500,8 +495,12 @@ def _check_axis(self, axis):
             raise ValueError("Duplicate axis: %s" % self.axis)
 
     def _check_if_input_shape_is_none(self, input_shape):
-        dim1, dim2 = input_shape[self.axis[0]], input_shape[self.axis[1]]
-        if dim1 is None or dim2 is None:
+        dims = [input_shape[i] for i in self.axis]
+
+        if len(input_shape) < 3:
+            raise ValueError("Expected input tensor to be at least 3D.")
+
+        if None in dims:
             raise ValueError(
                 """Axis {} of input tensor should have a defined dimension but
                 the layer received an input with shape {}.""".format(

tensorflow_addons/layers/tests/normalizations_test.py

@@ -368,25 +368,22 @@ def set_random_seed():
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 @pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
-def test_with_beta(dtype):
+@pytest.mark.parametrize("r_dim", [3, 4, 5, 6, 7])
+def test_with_beta(dtype, r_dim):
     set_random_seed()
 
-    shape = np.random.choice(range(1, 30), np.random.randint(2, 7), replace=True)
+    shape = np.random.choice(range(1, 30), int(r_dim), replace=True)
+    # shape = np.random.choice(range(1, 30), np.random.randint(2, 7), replace=True)
     inputs = np.random.random_sample(shape).astype(dtype)
+    print(shape)
 
-    if len(shape) == 2:
-        axis = channel_idx = 1
-
-    else:
-        axis = list(
-            np.random.choice(
-                range(1, len(shape) - 1),
-                np.random.randint(2, len(shape) - 1),
-                replace=False,
-            )
+    axis = list(
+        np.random.choice(
+            range(1, len(shape)), np.random.randint(1, len(shape) - 1), replace=False
         )
-        channel_idx = list(set(range(len(shape))) - set(axis) - set([0]))
-        channel_idx = int(np.random.choice(channel_idx, 1))
+    )
+    channel_idx = list(set(range(len(shape))) - set(axis) - set([0]))
+    channel_idx = int(np.random.choice(channel_idx, 1))
 
     frn = FilterResponseNormalization(
         beta_initializer="ones",
@@ -406,25 +403,20 @@ def test_with_beta(dtype):
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 @pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
-def test_with_gamma(dtype):
+@pytest.mark.parametrize("r_dim", [3, 4, 5, 6, 7])
+def test_with_gamma(dtype, r_dim):
     set_random_seed()
 
-    shape = np.random.choice(range(1, 30), np.random.randint(2, 7), replace=True)
+    shape = np.random.choice(range(1, 30), r_dim, replace=True)
     inputs = np.random.random_sample(shape).astype(dtype)
 
-    if len(shape) == 2:
-        axis = channel_idx = 1
-
-    else:
-        axis = list(
-            np.random.choice(
-                range(1, len(shape) - 1),
-                np.random.randint(2, len(shape) - 1),
-                replace=False,
-            )
+    axis = list(
+        np.random.choice(
+            range(1, len(shape)), np.random.randint(1, len(shape) - 1), replace=False
         )
-        channel_idx = list(set(range(len(shape))) - set(axis) - set([0]))
-        channel_idx = int(np.random.choice(channel_idx, 1))
+    )
+    channel_idx = list(set(range(len(shape))) - set(axis) - set([0]))
+    channel_idx = int(np.random.choice(channel_idx, 1))
 
     frn = FilterResponseNormalization(
         beta_initializer="zeros",
@@ -444,25 +436,20 @@ def test_with_gamma(dtype):
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 @pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
-def test_with_epsilon(dtype):
+@pytest.mark.parametrize("r_dim", [3, 4, 5, 6, 7])
+def test_with_epsilon(dtype, r_dim):
     set_random_seed()
 
-    shape = np.random.choice(range(1, 30), np.random.randint(2, 7), replace=True)
+    shape = np.random.choice(range(1, 30), r_dim, replace=True)
     inputs = np.random.random_sample(shape).astype(dtype)
 
-    if len(shape) == 2:
-        axis = channel_idx = 1
-
-    else:
-        axis = list(
-            np.random.choice(
-                range(1, len(shape) - 1),
-                np.random.randint(2, len(shape) - 1),
-                replace=False,
-            )
+    axis = list(
+        np.random.choice(
+            range(1, len(shape)), np.random.randint(1, len(shape) - 1), replace=False
         )
-        channel_idx = list(set(range(len(shape))) - set(axis) - set([0]))
-        channel_idx = int(np.random.choice(channel_idx, 1))
+    )
+    channel_idx = list(set(range(len(shape))) - set(axis) - set([0]))
+    channel_idx = int(np.random.choice(channel_idx, 1))
 
     frn = FilterResponseNormalization(
         beta_initializer=tf.keras.initializers.Constant(0.5),
@@ -489,27 +476,22 @@ def test_with_epsilon(dtype):
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 @pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
-def test_keras_model(dtype):
+@pytest.mark.parametrize("r_dim", [3, 4, 5, 6, 7])
+def test_keras_model(dtype, r_dim):
     set_random_seed()
 
-    shape = np.random.choice(range(1, 30), np.random.randint(2, 7), replace=True)
+    shape = np.random.choice(range(1, 30), r_dim, replace=True)
     random_inputs = np.random.random_sample(shape).astype(dtype)
     random_labels = np.random.randint(2, size=(shape[0],)).astype(dtype)
     input_layer = tf.keras.layers.Input(shape=tuple(shape[1:]))
 
-    if len(shape) == 2:
-        axis = channel_idx = 1
-
-    else:
-        axis = list(
-            np.random.choice(
-                range(1, len(shape) - 1),
-                np.random.randint(2, len(shape) - 1),
-                replace=False,
-            )
+    axis = list(
+        np.random.choice(
+            range(1, len(shape)), np.random.randint(1, len(shape) - 1), replace=False
         )
-        channel_idx = list(set(range(len(shape))) - set(axis) - set([0]))
-        channel_idx = int(np.random.choice(channel_idx, 1))
+    )
+    channel_idx = list(set(range(len(shape))) - set(axis) - set([0]))
+    channel_idx = int(np.random.choice(channel_idx, 1))
 
     frn = FilterResponseNormalization(
         beta_initializer="ones",
@@ -539,27 +521,22 @@ def test_serialization(dtype):
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 @pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
-def test_eps_grads(dtype):
+@pytest.mark.parametrize("r_dim", [3, 4, 5, 6, 7])
+def test_eps_grads(dtype, r_dim):
     set_random_seed()
 
-    shape = np.random.choice(range(1, 30), np.random.randint(2, 7), replace=True)
+    shape = np.random.choice(range(1, 30), r_dim, replace=True)
     random_inputs = np.random.random_sample(shape).astype(dtype)
     random_labels = np.random.randint(2, size=(shape[0],)).astype(dtype)
     input_layer = tf.keras.layers.Input(shape=tuple(shape[1:]))
 
-    if len(shape) == 2:
-        axis = channel_idx = 1
-
-    else:
-        axis = list(
-            np.random.choice(
-                range(1, len(shape) - 1),
-                np.random.randint(2, len(shape) - 1),
-                replace=False,
-            )
+    axis = list(
+        np.random.choice(
+            range(1, len(shape)), np.random.randint(1, len(shape) - 1), replace=False
         )
-        channel_idx = list(set(range(len(shape))) - set(axis) - set([0]))
-        channel_idx = int(np.random.choice(channel_idx, 1))
+    )
+    channel_idx = list(set(range(len(shape))) - set(axis) - set([0]))
+    channel_idx = int(np.random.choice(channel_idx, 1))
 
     frn = FilterResponseNormalization(
         beta_initializer="ones",