Squeeze and Excite block

keras_cv/layers/__init__.py

@@ -58,4 +58,5 @@
 from keras_cv.layers.preprocessing.solarization import Solarization
 from keras_cv.layers.regularization.drop_path import DropPath
 from keras_cv.layers.regularization.dropblock_2d import DropBlock2D
+from keras_cv.layers.regularization.squeeze_excite import SqueezeAndExciteBlock2D
 from keras_cv.layers.regularization.stochastic_depth import StochasticDepth

keras_cv/layers/regularization/__init__.py

@@ -14,4 +14,5 @@
 
 from keras_cv.layers.regularization.drop_path import DropPath
 from keras_cv.layers.regularization.dropblock_2d import DropBlock2D
+from keras_cv.layers.regularization.squeeze_excite import SqueezeAndExciteBlock2D
 from keras_cv.layers.regularization.stochastic_depth import StochasticDepth

keras_cv/layers/regularization/squeeze_excite.py

@@ -0,0 +1,73 @@
+# Copyright 2022 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tensorflow as tf
+from tensorflow.keras import layers
+
+
+@tf.keras.utils.register_keras_serializable(package="keras_cv")
+class SqueezeAndExciteBlock2D(layers.Layer):
+    """
+    Implements Squeeze and Excite block as in
+    [Squeeze-and-Excitation Networks](https://arxiv.org/pdf/1709.01507.pdf).
+
+    Args:
+            filters: Number of input and output filters. The number of input and
+                    output filters is same.
+            ratio: Ratio for bottleneck filters. Number of bottleneck filters =
+                    filters * ratio. Defaults to 0.25.
+    Usage:
+
+    ```python
+    # (...)
+    input = tf.ones((1, 5, 5, 16), dtype=tf.float32)
+    x = tf.keras.layers.Conv2D(16, (3, 3))(input)
+    output = keras_cv.layers.SqueezeAndExciteBlock(16)(x)
+    # (...)
+    ```
+    """
+
+    def __init__(self, filters, ratio=0.25, **kwargs):
+        super().__init__(**kwargs)
+
+        self.filters = filters
+
+        if ratio <= 0.0 or ratio >= 1.0:
+            raise ValueError(f"`ratio` should be a float between 0 and 1. Got {ratio}")
+
+        if filters <= 0 or not isinstance(filters, int):
+            raise ValueError(f"`filters` should be a positive integer. Got {filters}")
+
+        self.ratio = ratio
+        self.se_filters = int(self.filters * self.ratio)
+
+        self.ga_pool = layers.GlobalAveragePooling2D(keepdims=True)
+        self.squeeze_conv = layers.Conv2D(
+            self.se_filters,
+            (1, 1),
+            activation="relu",
+        )
+        self.excite_conv = layers.Conv2D(self.filters, (1, 1), activation="sigmoid")
+
+    def call(self, inputs, training=True):
+        x = self.ga_pool(inputs)  # x: (batch_size, 1, 1, filters)
+        x = self.squeeze_conv(x)  # x: (batch_size, 1, 1, se_filters)
+        x = self.excite_conv(x)  # x: (batch_size, 1, 1, filters)
+        x = tf.math.multiply(x, inputs)  # x: (batch_size, h, w, filters)
+        return x
+
+    def get_config(self):
+        config = {"filters": self.filters, "ratio": self.ratio}
+        base_config = super().get_config()
+        return dict(list(base_config.items()) + list(config.items()))

keras_cv/layers/regularization/squeeze_excite_test.py

@@ -0,0 +1,39 @@
+# Copyright 2022 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tensorflow as tf
+
+from keras_cv.layers import SqueezeAndExciteBlock2D
+
+
+class SqueezeAndExciteBlock2DTest(tf.test.TestCase):
+    def test_maintains_shape(self):
+        input_shape = (1, 4, 4, 8)
+        inputs = tf.random.uniform(input_shape)
+
+        layer = SqueezeAndExciteBlock2D(8, ratio=0.25)
+        outputs = layer(inputs)
+        self.assertEquals(inputs.shape, outputs.shape)
+
+    def test_raises_invalid_ratio_error(self):
+        with self.assertRaisesRegex(
+            ValueError, "`ratio` should be a float" " between 0 and 1. Got (.*?)"
+        ):
+            _ = SqueezeAndExciteBlock2D(8, ratio=1.1)
+
+    def test_raises_invalid_filters_error(self):
+        with self.assertRaisesRegex(
+            ValueError, "`filters` should be a positive" " integer. Got (.*?)"
+        ):
+            _ = SqueezeAndExciteBlock2D(-8.7)

keras_cv/layers/serialization_test.py

@@ -127,6 +127,11 @@ class SerializationTest(tf.test.TestCase, parameterized.TestCase):
             regularization.StochasticDepth,
             {"rate": 0.1},
         ),
+        (
+            "SqueezeAndExciteBlock2D",
+            regularization.SqueezeAndExciteBlock2D,
+            {"filters": 16, "ratio": 0.25},
+        ),
         (
             "DropPath",
             regularization.DropPath,