Add SpatialDropout1D/2D/3D layers

keras_core/layers/__init__.py

@@ -10,3 +10,6 @@ from keras_core.layers.regularization.activity_regularization import (
 from keras_core.layers.regularization.dropout import Dropout
 from keras_core.layers.regularization.gaussian_dropout import GaussianDropout
 from keras_core.layers.regularization.gaussian_noise import GaussianNoise
+from keras_core.layers.regularization.spatial_dropout import SpatialDropout1D
+from keras_core.layers.regularization.spatial_dropout import SpatialDropout2D
+from keras_core.layers.regularization.spatial_dropout import SpatialDropout3D

keras_core/layers/regularization/spatial_dropout.py

@@ -0,0 +1,204 @@
+from keras_core import backend
+from keras_core import operations as ops
+from keras_core.api_export import keras_core_export
+from keras_core.layers.input_spec import InputSpec
+from keras_core.layers.regularization.dropout import Dropout
+
+
+class BaseSpatialDropout(Dropout):
+    def __init__(self, rate, seed=None, name=None, dtype=None):
+        super().__init__(rate, seed=seed, name=name, dtype=dtype)
+
+    def call(self, inputs, training=False):
+        if training and self.rate > 0:
+            return backend.random.dropout(
+                inputs,
+                self.rate,
+                noise_shape=self._get_noise_shape(inputs),
+                seed=self.seed_generator,
+            )
+        return inputs
+
+    def get_config(self):
+        return {
+            "rate": self.rate,
+            "seed": self.seed,
+            "name": self.name,
+            "dtype": self.dtype,
+        }
+
+
+@keras_core_export("keras_core.layers.SpatialDropout1D")
+class SpatialDropout1D(BaseSpatialDropout):
+    """Spatial 1D version of Dropout.
+
+    This layer performs the same function as Dropout, however, it drops
+    entire 1D feature maps instead of individual elements. If adjacent frames
+    within feature maps are strongly correlated (as is normally the case in
+    early convolution layers) then regular dropout will not regularize the
+    activations and will otherwise just result in an effective learning rate
+    decrease. In this case, `SpatialDropout1D` will help promote independence
+    between feature maps and should be used instead.
+
+    Args:
+        rate: Float between 0 and 1. Fraction of the input units to drop.
+
+    Call arguments:
+        inputs: A 3D tensor.
+        training: Python boolean indicating whether the layer
+            should behave in training mode (applying dropout)
+            or in inference mode (pass-through).
+
+    Input shape:
+        3D tensor with shape: `(samples, timesteps, channels)`
+
+    Output shape: Same as input.
+
+    Reference:
+
+    - [Tompson et al., 2014](https://arxiv.org/abs/1411.4280)
+    """
+
+    def __init__(self, rate, seed=None, name=None, dtype=None):
+        super().__init__(rate, seed=seed, name=name, dtype=dtype)
+        self.input_spec = InputSpec(ndim=3)
+
+    def _get_noise_shape(self, inputs):
+        input_shape = ops.shape(inputs)
+        return (input_shape[0], 1, input_shape[2])
+
+
+@keras_core_export("keras_core.layers.SpatialDropout2D")
+class SpatialDropout2D(BaseSpatialDropout):
+    """Spatial 2D version of Dropout.
+
+    This version performs the same function as Dropout, however, it drops
+    entire 2D feature maps instead of individual elements. If adjacent pixels
+    within feature maps are strongly correlated (as is normally the case in
+    early convolution layers) then regular dropout will not regularize the
+    activations and will otherwise just result in an effective learning rate
+    decrease. In this case, `SpatialDropout2D` will help promote independence
+    between feature maps and should be used instead.
+
+    Args:
+        rate: Float between 0 and 1. Fraction of the input units to drop.
+        data_format: `"channels_first"` or `"channels_last"`.
+            In `"channels_first"` mode, the channels dimension (the depth)
+            is at index 1, in `"channels_last"` mode is it at index 3.
+            It defaults to the `image_data_format` value found in your
+            Keras config file at `~/.keras/keras.json`.
+            If you never set it, then it will be `"channels_last"`.
+
+    Call arguments:
+        inputs: A 4D tensor.
+        training: Python boolean indicating whether the layer
+            should behave in training mode (applying dropout)
+            or in inference mode (pass-through).
+
+    Input shape:
+        4D tensor with shape: `(samples, channels, rows, cols)` if
+            data_format='channels_first'
+        or 4D tensor with shape: `(samples, rows, cols, channels)` if
+            data_format='channels_last'.
+
+    Output shape: Same as input.
+
+    Reference:
+
+    - [Tompson et al., 2014](https://arxiv.org/abs/1411.4280)
+    """
+
+    def __init__(
+        self, rate, data_format=None, seed=None, name=None, dtype=None
+    ):
+        super().__init__(rate, seed=seed, name=name, dtype=dtype)
+        data_format = data_format or backend.image_data_format()
+        if data_format not in {"channels_last", "channels_first"}:
+            raise ValueError(
+                '`data_format` must be "channels_last" or "channels_first". '
+                f"Received: data_format={data_format}."
+            )
+        self.data_format = data_format
+        self.input_spec = InputSpec(ndim=4)
+
+    def _get_noise_shape(self, inputs):
+        input_shape = ops.shape(inputs)
+        if self.data_format == "channels_first":
+            return (input_shape[0], input_shape[1], 1, 1)
+        elif self.data_format == "channels_last":
+            return (input_shape[0], 1, 1, input_shape[3])
+
+    def get_config(self):
+        base_config = super().get_config()
+        config = {
+            "data_format": self.data_format,
+        }
+        return {**base_config, **config}
+
+
+@keras_core_export("keras_core.layers.SpatialDropout3D")
+class SpatialDropout3D(BaseSpatialDropout):
+    """Spatial 3D version of Dropout.
+
+    This version performs the same function as Dropout, however, it drops
+    entire 3D feature maps instead of individual elements. If adjacent voxels
+    within feature maps are strongly correlated (as is normally the case in
+    early convolution layers) then regular dropout will not regularize the
+    activations and will otherwise just result in an effective learning rate
+    decrease. In this case, SpatialDropout3D will help promote independence
+    between feature maps and should be used instead.
+
+    Args:
+        rate: Float between 0 and 1. Fraction of the input units to drop.
+        data_format: `"channels_first"` or `"channels_last"`.
+            In `"channels_first"` mode, the channels dimension (the depth)
+            is at index 1, in `"channels_last"` mode is it at index 4.
+            It defaults to the `image_data_format` value found in your
+            Keras config file at `~/.keras/keras.json`.
+            If you never set it, then it will be `"channels_last"`.
+
+    Call arguments:
+        inputs: A 5D tensor.
+        training: Python boolean indicating whether the layer
+                should behave in training mode (applying dropout)
+                or in inference mode (pass-through).
+
+    Input shape:
+        5D tensor with shape: `(samples, channels, dim1, dim2, dim3)` if
+            data_format='channels_first'
+        or 5D tensor with shape: `(samples, dim1, dim2, dim3, channels)` if
+            data_format='channels_last'.
+
+    Output shape: Same as input.
+
+    Reference:
+
+    - [Tompson et al., 2014](https://arxiv.org/abs/1411.4280)
+    """
+
+    def __init__(
+        self, rate, data_format=None, seed=None, name=None, dtype=None
+    ):
+        super().__init__(rate, seed=seed, name=name, dtype=dtype)
+        data_format = data_format or backend.image_data_format()
+        if data_format not in {"channels_last", "channels_first"}:
+            raise ValueError(
+                '`data_format` must be "channels_last" or "channels_first". '
+                f"Received: data_format={data_format}."
+            )
+        self.data_format = data_format
+        self.input_spec = InputSpec(ndim=5)
+
+    def _get_noise_shape(self, inputs):
+        input_shape = ops.shape(inputs)
+        if self.data_format == "channels_first":
+            return (input_shape[0], input_shape[1], 1, 1, 1)
+        elif self.data_format == "channels_last":
+            return (input_shape[0], 1, 1, 1, input_shape[4])
+
+    def get_config(self):
+        base_config = super().get_config()
+        config = {
+            "data_format": self.data_format,
+        }
+        return {**base_config, **config}

keras_core/layers/regularization/spatial_dropout_test.py

@@ -0,0 +1,98 @@
+import numpy as np
+import pytest
+
+from keras_core import backend
+from keras_core import layers
+from keras_core.testing import test_case
+
+
+class SpatialDropoutTest(test_case.TestCase):
+    def test_spatial_dropout_1d(self):
+        self.run_layer_test(
+            layers.SpatialDropout1D,
+            init_kwargs={"rate": 0.5},
+            call_kwargs={"training": True},
+            input_shape=(2, 3, 4),
+        )
+
+        self.run_layer_test(
+            layers.SpatialDropout1D,
+            init_kwargs={"rate": 0.5},
+            call_kwargs={"training": False},
+            input_shape=(2, 3, 4),
+        )
+
+    def test_spatial_dropout_2d(self):
+        self.run_layer_test(
+            layers.SpatialDropout2D,
+            init_kwargs={"rate": 0.5},
+            call_kwargs={"training": True},
+            input_shape=(2, 3, 4, 5),
+        )
+
+        self.run_layer_test(
+            layers.SpatialDropout2D,
+            init_kwargs={"rate": 0.5, "data_format": "channels_first"},
+            call_kwargs={"training": True},
+            input_shape=(2, 3, 4, 5),
+        )
+
+    def test_spatial_dropout_3d(self):
+        self.run_layer_test(
+            layers.SpatialDropout3D,
+            init_kwargs={"rate": 0.5},
+            call_kwargs={"training": True},
+            input_shape=(2, 3, 4, 4, 5),
+        )
+
+        self.run_layer_test(
+            layers.SpatialDropout3D,
+            init_kwargs={"rate": 0.5, "data_format": "channels_first"},
+            call_kwargs={"training": True},
+            input_shape=(2, 3, 4, 4, 5),
+        )
+
+    @pytest.mark.skipif(
+        not backend.DYNAMIC_SHAPES_OK,
+        reason="Backend does not support dynamic shapes",
+    )
+    def test_spatial_dropout_1D_dynamic(self):
+        inputs = layers.Input((3, 2))
+        layer = layers.SpatialDropout1D(0.5)
+        layer(inputs, training=True)
+
+    def test_spatial_dropout_1D_correctness(self):
+        inputs = np.ones((10, 3, 10))
+        layer = layers.SpatialDropout1D(0.5)
+        outputs = layer(inputs, training=True)
+        self.assertAllClose(outputs[:, 0, :], outputs[:, 1, :])
+
+    @pytest.mark.skipif(
+        not backend.DYNAMIC_SHAPES_OK,
+        reason="Backend does not support dynamic shapes",
+    )
+    def test_spatial_dropout_2D_dynamic(self):
+        inputs = layers.Input((3, 2, 4))
+        layer = layers.SpatialDropout2D(0.5)
+        layer(inputs, training=True)
+
+    def test_spatial_dropout_2D_correctness(self):
+        inputs = np.ones((10, 3, 3, 10))
+        layer = layers.SpatialDropout2D(0.5)
+        outputs = layer(inputs, training=True)
+        self.assertAllClose(outputs[:, 0, 0, :], outputs[:, 1, 1, :])
+
+    @pytest.mark.skipif(
+        not backend.DYNAMIC_SHAPES_OK,
+        reason="Backend does not support dynamic shapes",
+    )
+    def test_spatial_dropout_3D_dynamic(self):
+        inputs = layers.Input((3, 2, 4, 2))
+        layer = layers.SpatialDropout3D(0.5)
+        layer(inputs, training=True)
+
+    def test_spatial_dropout_3D_correctness(self):
+        inputs = np.ones((10, 3, 3, 3, 10))
+        layer = layers.SpatialDropout3D(0.5)
+        outputs = layer(inputs, training=True)
+        self.assertAllClose(outputs[:, 0, 0, 0, :], outputs[:, 1, 1, 1, :])