Add demo

README.md

@@ -13,6 +13,4 @@ pip install keras-drop-block
 
 ## Usage
 
-```python
-
-```
+See [fashion mnist demo](./demo/mnist.py).

demo/mnist.py

@@ -0,0 +1,86 @@
+import keras
+import keras.backend as K
+import numpy as np
+from keras.datasets import fashion_mnist
+from keras_drop_block import DropBlock2D
+
+
+(x_train, y_train), (x_test, y_test) = fashion_mnist.load_data()
+
+x_train = np.expand_dims(x_train.astype(K.floatx()) / 255, axis=-1)
+x_test = np.expand_dims(x_test.astype(K.floatx()) / 255, axis=-1)
+
+y_train, y_test = np.expand_dims(y_train, axis=-1), np.expand_dims(y_test, axis=-1)
+
+train_num = round(x_train.shape[0] * 0.9)
+x_train, x_valid = x_train[:train_num, ...], x_train[train_num:, ...]
+y_train, y_valid = y_train[:train_num, ...], y_train[train_num:, ...]
+
+
+def get_dropout_model():
+    model = keras.models.Sequential()
+    model.add(keras.layers.Dropout(input_shape=(28, 28, 1), rate=0.3, name='Input-Dropout'))
+    model.add(keras.layers.Conv2D(filters=64, kernel_size=3, activation='relu', padding='same', name='Conv-1'))
+    model.add(keras.layers.MaxPool2D(pool_size=2, name='Pool-1'))
+    model.add(keras.layers.Dropout(rate=0.2, name='Dropout-1'))
+    model.add(keras.layers.Conv2D(filters=32, kernel_size=3, activation='relu', padding='same', name='Conv-2'))
+    model.add(keras.layers.MaxPool2D(pool_size=2, name='Pool-2'))
+    model.add(keras.layers.Dropout(rate=0.2, name='Dropout-2'))
+    model.add(keras.layers.Flatten(name='Flatten'))
+    model.add(keras.layers.Dense(units=256, activation='relu', name='Dense'))
+    model.add(keras.layers.Dropout(rate=0.2, name='Dense-Dropout'))
+    model.add(keras.layers.Dense(units=10, activation='softmax', name='Softmax'))
+    model.compile(
+        optimizer='adam',
+        loss='sparse_categorical_crossentropy',
+        metrics=['accuracy'],
+    )
+    return model
+
+
+dropout_model = get_dropout_model()
+dropout_model.summary()
+dropout_model.fit(
+    x=x_train,
+    y=y_train,
+    epochs=10,
+    validation_data=(x_valid, y_valid),
+    callbacks=[keras.callbacks.EarlyStopping(monitor='val_acc', patience=2)]
+)
+dropout_score = dropout_model.evaluate(x_test, y_test)
+print('Score of dropout:\t%.4f' % dropout_score[1])
+
+
+def get_drop_block_model():
+    model = keras.models.Sequential()
+    model.add(DropBlock2D(input_shape=(28, 28, 1), block_size=7, keep_prob=0.8, name='Input-Dropout'))
+    model.add(keras.layers.Conv2D(filters=64, kernel_size=3, activation='relu', padding='same', name='Conv-1'))
+    model.add(keras.layers.MaxPool2D(pool_size=2, name='Pool-1'))
+    model.add(DropBlock2D(block_size=5, keep_prob=0.8, name='Dropout-1'))
+    model.add(keras.layers.Conv2D(filters=32, kernel_size=3, activation='relu', padding='same', name='Conv-2'))
+    model.add(keras.layers.MaxPool2D(pool_size=2, name='Pool-2'))
+    model.add(DropBlock2D(block_size=3, keep_prob=0.8, name='Dropout-2'))
+    model.add(keras.layers.Flatten(name='Flatten'))
+    model.add(keras.layers.Dense(units=256, activation='relu', name='Dense'))
+    model.add(keras.layers.Dropout(rate=0.2, name='Dense-Dropout'))
+    model.add(keras.layers.Dense(units=10, activation='softmax', name='Softmax'))
+    model.compile(
+        optimizer='adam',
+        loss='sparse_categorical_crossentropy',
+        metrics=['accuracy'],
+    )
+    return model
+
+
+drop_block_model = get_drop_block_model()
+drop_block_model.summary()
+drop_block_model.fit(
+    x=x_train,
+    y=y_train,
+    epochs=10,
+    validation_data=(x_valid, y_valid),
+    callbacks=[keras.callbacks.EarlyStopping(monitor='val_acc', patience=2)]
+)
+drop_block_score = drop_block_model.evaluate(x_test, y_test)
+print('Score of dropout:\t%.4f' % dropout_score[1])
+print('Score of DropBlock:\t%.4f' % drop_block_score[1])

keras_drop_block/__init__.py

@@ -1 +1 @@
-from .drop_block import DropBlock2D
+from .drop_block import DropBlock1D, DropBlock2D

keras_drop_block/drop_block.py

@@ -2,31 +2,110 @@
 import keras.backend as K
 
 
+class DropBlock1D(keras.layers.Layer):
+    """See: https://arxiv.org/pdf/1810.12890.pdf"""
+
+    def __init__(self,
+                 block_size,
+                 keep_prob,
+                 sync_channels=False,
+                 data_format=None,
+                 **kwargs):
+        """Initialize the layer.
+
+        :param block_size: Size for each mask block.
+        :param keep_prob: Probability of keeping the original feature.
+        :param sync_channels: Whether to use the same dropout for all channels.
+        :param data_format: 'channels_first' or 'channels_last' (default).
+        :param kwargs: Arguments for parent class.
+        """
+        super(DropBlock1D, self).__init__(**kwargs)
+        self.block_size = block_size
+        self.keep_prob = keep_prob
+        self.sync_channels = sync_channels
+        self.data_format = K.normalize_data_format(data_format)
+        self.input_spec = keras.engine.base_layer.InputSpec(ndim=3)
+        self.supports_masking = True
+
+    def get_config(self):
+        config = {'block_size': self.block_size,
+                  'keep_prob': self.keep_prob,
+                  'sync_channels': self.sync_channels,
+                  'data_format': self.data_format}
+        base_config = super(DropBlock1D, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    def compute_mask(self, inputs, mask=None):
+        return mask
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+    def _get_gamma(self, feature_dim):
+        """Get the number of activation units to drop"""
+        feature_dim = K.cast(feature_dim, K.floatx())
+        block_size = K.constant(self.block_size, dtype=K.floatx())
+        return ((1.0 - self.keep_prob) / block_size) * (feature_dim / (feature_dim - block_size + 1.0))
+
+    def _compute_drop_mask(self, shape):
+        mask = K.random_binomial(shape, p=self._get_gamma(shape[1]))
+        mask = keras.layers.MaxPool1D(
+            pool_size=self.block_size,
+            padding='same',
+            strides=1,
+            data_format='channels_last',
+        )(mask)
+        return 1.0 - mask
+
+    def call(self, inputs, training=None):
+
+        def dropped_inputs():
+            outputs = inputs
+            if self.data_format == 'channels_first':
+                outputs = K.permute_dimensions(outputs, [0, 2, 1])
+            shape = K.shape(outputs)
+            if self.sync_channels:
+                mask = self._compute_drop_mask([shape[0], shape[1], 1])
+            else:
+                mask = self._compute_drop_mask(shape)
+            outputs = outputs * mask *\
+                (K.cast(K.prod(shape), dtype=K.floatx()) / K.sum(mask))
+            if self.data_format == 'channels_first':
+                outputs = K.permute_dimensions(outputs, [0, 2, 1])
+            return outputs
+
+        return K.in_train_phase(dropped_inputs, inputs, training=training)
+
+
 class DropBlock2D(keras.layers.Layer):
     """See: https://arxiv.org/pdf/1810.12890.pdf"""
 
     def __init__(self,
                  block_size,
                  keep_prob,
+                 sync_channels=False,
                  data_format=None,
                  **kwargs):
         """Initialize the layer.
 
         :param block_size: Size for each mask block.
         :param keep_prob: Probability of keeping the original feature.
+        :param sync_channels: Whether to use the same dropout for all channels.
         :param data_format: 'channels_first' or 'channels_last' (default).
         :param kwargs: Arguments for parent class.
         """
         super(DropBlock2D, self).__init__(**kwargs)
         self.block_size = block_size
         self.keep_prob = keep_prob
+        self.sync_channels = sync_channels
         self.data_format = K.normalize_data_format(data_format)
         self.input_spec = keras.engine.base_layer.InputSpec(ndim=4)
         self.supports_masking = True
 
     def get_config(self):
         config = {'block_size': self.block_size,
                   'keep_prob': self.keep_prob,
+                  'sync_channels': self.sync_channels,
                   'data_format': self.data_format}
         base_config = super(DropBlock2D, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
@@ -44,7 +123,7 @@ def _get_gamma(self, height, width):
         return ((1.0 - self.keep_prob) / (block_size ** 2)) *\
                (height * width / ((height - block_size + 1.0) * (width - block_size + 1.0)))
 
-    def _compute_mask(self, shape):
+    def _compute_drop_mask(self, shape):
         mask = K.random_binomial(shape, p=self._get_gamma(shape[1], shape[2]))
         mask = keras.layers.MaxPool2D(
             pool_size=(self.block_size, self.block_size),
@@ -61,7 +140,10 @@ def dropped_inputs():
             if self.data_format == 'channels_first':
                 outputs = K.permute_dimensions(outputs, [0, 2, 3, 1])
             shape = K.shape(outputs)
-            mask = self._compute_mask(shape)
+            if self.sync_channels:
+                mask = self._compute_drop_mask([shape[0], shape[1], shape[2], 1])
+            else:
+                mask = self._compute_drop_mask(shape)
             outputs = outputs * mask *\
                 (K.cast(K.prod(shape), dtype=K.floatx()) / K.sum(mask))
             if self.data_format == 'channels_first':

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='keras-drop-block',
-    version='0.1.0',
+    version='0.2.0',
     packages=find_packages(),
     url='https://github.com/CyberZHG/keras-drop-block',
     license='MIT',

tests/test_drop_block_1d.py

@@ -0,0 +1,109 @@
+import os
+import random
+import tempfile
+import unittest
+import keras
+import numpy as np
+from keras_drop_block import DropBlock1D
+
+
+class TestDropBlock1D(unittest.TestCase):
+
+    def test_training(self):
+        input_layer = keras.layers.Input(shape=(100, 3))
+        drop_block_layer = DropBlock1D(block_size=3, keep_prob=0.7)(input_layer)
+        model = keras.models.Model(inputs=input_layer, outputs=drop_block_layer)
+        model.compile(optimizer='adam', loss='mse', metrics={})
+        model_path = os.path.join(tempfile.gettempdir(), 'keras_drop_block_%f.h5' % random.random())
+        model.save(model_path)
+        model = keras.models.load_model(
+            model_path,
+            custom_objects={'DropBlock1D': DropBlock1D},
+        )
+        model.summary()
+        inputs = np.ones((1, 100, 3))
+        outputs = model.predict(inputs)
+        self.assertTrue(np.allclose(inputs, outputs))
+
+        input_layer = keras.layers.Input(shape=(3, 100))
+        drop_block_layer = DropBlock1D(block_size=3, keep_prob=0.7, data_format='channels_first')(input_layer)
+        model = keras.models.Model(inputs=input_layer, outputs=drop_block_layer)
+        model.compile(optimizer='adam', loss='mse', metrics={})
+        model_path = os.path.join(tempfile.gettempdir(), 'keras_drop_block_%f.h5' % random.random())
+        model.save(model_path)
+        model = keras.models.load_model(
+            model_path,
+            custom_objects={'DropBlock1D': DropBlock1D},
+        )
+        model.summary()
+        inputs = np.ones((1, 3, 100))
+        outputs = model.predict(inputs)
+        self.assertTrue(np.allclose(inputs, outputs))
+
+    def test_mask_shape(self):
+        input_layer = keras.layers.Input(shape=(100, 3))
+        drop_block_layer = keras.layers.Lambda(
+            lambda x: DropBlock1D(block_size=3, keep_prob=0.7)(x, training=True),
+        )(input_layer)
+        model = keras.models.Model(inputs=input_layer, outputs=drop_block_layer)
+        model.compile(optimizer='adam', loss='mse', metrics={})
+        model_path = os.path.join(tempfile.gettempdir(), 'keras_drop_block_%f.h5' % random.random())
+        model.save(model_path)
+        model = keras.models.load_model(
+            model_path,
+            custom_objects={'DropBlock1D': DropBlock1D},
+        )
+        model.summary()
+        inputs = np.ones((1, 100, 3))
+        outputs = model.predict(inputs)
+        for i in range(3):
+            print((outputs[0, :, i] > 0.0).astype(dtype='int32'))
+        inputs = np.ones((1000, 100, 3))
+        outputs = model.predict(inputs)
+        keep_prob = 1.0 * np.sum(outputs > 0.0) / np.prod(np.shape(outputs))
+        self.assertTrue(0.6 < keep_prob < 0.8, keep_prob)
+
+        input_layer = keras.layers.Input(shape=(3, 100))
+        drop_block_layer = keras.layers.Lambda(
+            lambda x: DropBlock1D(block_size=3, keep_prob=0.7, data_format='channels_first')(x, training=True),
+        )(input_layer)
+        model = keras.models.Model(inputs=input_layer, outputs=drop_block_layer)
+        model.compile(optimizer='adam', loss='mse', metrics={})
+        model_path = os.path.join(tempfile.gettempdir(), 'keras_drop_block_%f.h5' % random.random())
+        model.save(model_path)
+        model = keras.models.load_model(
+            model_path,
+            custom_objects={'DropBlock1D': DropBlock1D},
+        )
+        model.summary()
+        inputs = np.ones((1, 3, 100))
+        outputs = model.predict(inputs)
+        for i in range(3):
+            print((outputs[0, i, :] > 0.0).astype(dtype='int32'))
+        inputs = np.ones((1000, 3, 100))
+        outputs = model.predict(inputs)
+        keep_prob = 1.0 * np.sum(outputs > 0.0) / np.prod(np.shape(outputs))
+        self.assertTrue(0.6 < keep_prob < 0.8, keep_prob)
+
+    def test_sync_channels(self):
+        input_layer = keras.layers.Input(shape=(100, 3))
+        drop_block_layer = keras.layers.Lambda(
+            lambda x: DropBlock1D(block_size=3, keep_prob=0.7, sync_channels=True)(x, training=True),
+        )(input_layer)
+        model = keras.models.Model(inputs=input_layer, outputs=drop_block_layer)
+        model.compile(optimizer='adam', loss='mse', metrics={})
+        model_path = os.path.join(tempfile.gettempdir(), 'keras_drop_block_%f.h5' % random.random())
+        model.save(model_path)
+        model = keras.models.load_model(
+            model_path,
+            custom_objects={'DropBlock1D': DropBlock1D},
+        )
+        model.summary()
+        inputs = np.ones((1, 100, 3))
+        outputs = model.predict(inputs)
+        for i in range(1, 3):
+            self.assertTrue(np.allclose(outputs[0, :, 0], outputs[0, :, i]))
+        inputs = np.ones((1000, 100, 3))
+        outputs = model.predict(inputs)
+        keep_prob = 1.0 * np.sum(outputs > 0.0) / np.prod(np.shape(outputs))
+        self.assertTrue(0.6 < keep_prob < 0.8, keep_prob)

tests/test_drop_block_2d.py

@@ -84,3 +84,26 @@ def test_mask_shape(self):
         outputs = model.predict(inputs)
         keep_prob = 1.0 * np.sum(outputs > 0.0) / np.prod(np.shape(outputs))
         self.assertTrue(0.6 < keep_prob < 0.8, keep_prob)
+
+    def test_sync_channels(self):
+        input_layer = keras.layers.Input(shape=(10, 10, 3))
+        drop_block_layer = keras.layers.Lambda(
+            lambda x: DropBlock2D(block_size=3, keep_prob=0.7, sync_channels=True)(x, training=True),
+        )(input_layer)
+        model = keras.models.Model(inputs=input_layer, outputs=drop_block_layer)
+        model.compile(optimizer='adam', loss='mse', metrics={})
+        model_path = os.path.join(tempfile.gettempdir(), 'keras_drop_block_%f.h5' % random.random())
+        model.save(model_path)
+        model = keras.models.load_model(
+            model_path,
+            custom_objects={'DropBlock2D': DropBlock2D},
+        )
+        model.summary()
+        inputs = np.ones((1, 10, 10, 3))
+        outputs = model.predict(inputs)
+        for i in range(1, 3):
+            self.assertTrue(np.allclose(outputs[0, :, :, 0], outputs[0, :, :, i]))
+        inputs = np.ones((1000, 10, 10, 3))
+        outputs = model.predict(inputs)
+        keep_prob = 1.0 * np.sum(outputs > 0.0) / np.prod(np.shape(outputs))
+        self.assertTrue(0.6 < keep_prob < 0.8, keep_prob)