diff --git a/keras_core/layers/__init__.py b/keras_core/layers/__init__.py index a6030e29b..4f2c3a2eb 100644 --- a/keras_core/layers/__init__.py +++ b/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 diff --git a/keras_core/layers/regularization/spatial_dropout.py b/keras_core/layers/regularization/spatial_dropout.py new file mode 100644 index 000000000..6068e4d0f --- /dev/null +++ b/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} diff --git a/keras_core/layers/regularization/spatial_dropout_test.py b/keras_core/layers/regularization/spatial_dropout_test.py new file mode 100644 index 000000000..dca0cfe02 --- /dev/null +++ b/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, :])