Introduce DropBlock2D regularization layer.

[sebastian-sz]

DropBlock2d is a regularization layer (similar to Dropout), more suitable for Convolutional networks.

Linked Issue: #137

This is mostly a Keras Layer wrapper around original implementation.

[LukeWood]

Thanks for the PR! Going forward we will be using the BaseImageAugmentationLayer. I'll have an example ready for that shortly... had to finish the base layer.

keras-team/keras#16161

[sebastian-sz]

@LukeWood Can we discuss it? I did not want to use BaseImageAugmentationLayer because DropBlock2D is not performing image augmentation - it's not meant to be used before model for modyfing input, but inside model to modify the feature maps.

Also the functionalities of BaseImageAugmentationLayer, like vectorized map, access to label modifications etc. are not necessary here, as the above implementation is vectorized and only cares about output from the previous layer.
I think it is much closer to Dropout, than to other preprocessing layers.

I do not want to break layer consistency in this repo, but as I said I'm not sure if BaseImageAugmentationLayer is necessary here?

[bhack]

I don't know if It makes sense to inherit from base_layer.BaseRandomLayer https://github.com/keras-team/keras/blob/master/keras/layers/regularization/dropout.py#L26

[sebastian-sz]

@bhack would you also say that the layer should inherit from BaseImageAugmentationLayer?

[bhack]

I meant that we could evaluate if this layer could inherit from BaseRandomLayer as the linked dropout.

[LukeWood]

@LukeWood Can we discuss it? I did not want to use BaseImageAugmentationLayer because DropBlock2D is not performing image augmentation - it's not meant to be used before model for modyfing input, but inside model to modify the feature maps.

Also the functionalities of BaseImageAugmentationLayer, like vectorized map, access to label modifications etc. are not necessary here, as the above implementation is vectorized and only cares about output from the previous layer.
I think it is much closer to Dropout, than to other preprocessing layers.

I do not want to break layer consistency in this repo, but as I said I'm not sure if BaseImageAugmentationLayer is necessary here?

Ah I see I did not realize how this was meant to be used. Let me take a look and review then.

[sebastian-sz]

@LukeWood thanks!

[LukeWood]

So far looking pretty good, some minor comments but overall great layer!

[LukeWood]

We might need to make this a tf.variable internally so it doesn't get baked into the graph to support a keep_probability schedule callback. Not sure yet, no action item needed on your side. Still doing some digging here.

[sebastian-sz]

Yes I was also thinking about the Schedule Callback. Something like

class DropBlockScheduleCallback(tf.keras.callbacks.Callback):
    def __init__(self, total_steps: int):
        super(DropBlockScheduleCallback, self).__init__()
        self._total_steps = total_steps
        self._current_step = 0

    def on_train_batch_end(self):
        self._current_step += 1
        for layer in self.model.layers:
            if isinstance(layer, DropBlock2D):
                new_keep_probability = self._calc_new_keep_probability(
                    layer._keep_probability  # TODO: public
                )
                layer._keep_probability.assign(new_keep_probability)

    def _calc_new_keep_probability(self, current_keep_probability):
        current_ratio = self._current_step / self._total_steps
        return 1 - current_ratio * (1 - current_keep_probability)

Although I'm not sure:

1. is there any other way to get total_steps without having them need to be explicitly provided.
2. is iterating over all model layers for each batch not too slow. Maybe they could be saved on_train_begin somewhere and used as a reference on_batch_end?

If I'm thinking in the right direction, I could add this in this PR. Or this could be a separate issue. Overall, please let me know.

[lucasdavid]

I don't know if this is easy to implement, but it would be very cool if could accept lambdas or Schedules:

global_step = tf.Variable(0, trainable=False)
db = DropBlock(keep_probability=lambda: 1/step)
# or ...
first_keep_prob = 1.0
db = DropBlock(
  keep_probability=tf.keras.optimizers.schedules.InverseTimeDecay(
    first_keep_prob,
    decay_steps=steps_per_epoch*epochs,
    decay_rate=0.96))

[sebastian-sz]

Yeah, I'm not sure what would be optimal solution for this parameter's scheduling. I would probably add it in a separate PR (+ separate issue), as there is still some work regarding the basic functionality.

[LukeWood]

I'll discuss this with the team and let you know. It would be good to get this sorted out because we actually will have the same problem with a few other components.

[sebastian-sz]

@qlzh727 @LukeWood
There is an issue, where if dropblock_size (height or width) is equal or bigger than input feature map, the layer will output nan values.
What is your take on validating the input size in the call method?

def call(...):
(...)
    tf.Assert(tf.less(self._dropblock_height, height), data=[''])
    tf.Assert(tf.less(self._dropblock_width, width), data=[''])

[qlzh727]

@qlzh727 @LukeWood There is an issue, where if dropblock_size (height or width) is equal or bigger than input feature map, the layer will output nan values. What is your take on validating the input size in the call method?

def call(...):
(...)
    tf.Assert(tf.less(self._dropblock_height, height), data=[''])
    tf.Assert(tf.less(self._dropblock_width, width), data=[''])

I think we should add such validation logic so that the mask block size is never bigger than the actual input.

[bhack]

I think we should add such validation logic so that the mask block size is never bigger than the actual input.

Isn't an user directed controlled parameter?
In other cases we are not protecting a misusing when we think that it is a caller ownership like in #144 (comment)

[qlzh727]

Thanks for the update.

[qlzh727]

Thanks for the PR.

[LukeWood]

Alright! This looks pretty great @sebastian-sz. Let me discuss the scheduling issue and then we can merge!

[lucasdavid]

Saving a model that contains this layer and re-loading it into memory is raising an "unexpected argument trainable" error. Snippet:

import tensorflow as tf
from keras import layers
model = tf.keras.Sequential([
  layers.Conv2D(32, 3)
  DropBlock2D(0.1, 16)
])
model.save('saved_model')
tf.keras.models.load_model('saved_model')

Adding **kwargs in the layer's initializer and passing it to the superclass fixes it:

@tf.keras.utils.register_keras_serializable('layers')
class DropBlock2D(BaseRandomLayer):
  def __init__(
      self,
      dropout_rate,
      dropblock_size,
      data_format=None,
      seed=None,
      name=None,
      **kwargs,
  ):
    super().__init__(seed=seed, name=name, force_generator=True, **kwargs)

[bhack]

Yes, thanks I think that more in general we need to add serialization tests at some point