Repository: tensorflow/addons

Issue:
Feature Request : Stochastic Depth/ResDrop
**Describe the feature and the current behavior/state.**
Stochastic depth is a regularization technique to train very deep residual networks (the authors for example train ResNets of up to 1202 layers). Particularly, it allows training of, on average, more shallow networks, that retain their full depth at inference time.

Previous attempts (https://github.com/tensorflow/addons/issues/626) seem to have been unnecessarily convoluted. This implemenation would be a simple layer, attached to the end of a residual branch (as suggested here https://github.com/tensorflow/tensorflow/issues/8817#issuecomment-290788333, which is also the way it is described in the Shake-Drop paper: https://arxiv.org/abs/1802.02375)

Noteably this means that, if momentum is being used, the dropped layers will still receive small updates due to their historic gradients (similar to Dropout).

**Relevant information**
- Are you willing to contribute it (yes/no): Yes
- Are you willing to maintain it going forward? (yes/no): Yes
- Is there a relevant academic paper? (if so, where): Yes (https://arxiv.org/abs/1603.09382)
- Is there already an implementation in another framework? (if so, where): The authors original pytorch implementation is available here: https://github.com/yueatsprograms/Stochastic_Depth
- Was it part of tf.contrib? (if so, where): No

**Which API type would this fall under (layer, metric, optimizer, etc.)**
layer
**Who will benefit with this feature?**
Anybody who wants to use stochastic depth to train deeper ResNets, or who wants to recreate the EfficientNet architecture (https://arxiv.org/abs/1905.11946). Anyone who wants to add Shake-Drop (https://arxiv.org/abs/1802.02375) to their network, can use this as base.
**Any other info.**
It is important to note that Stochastic Depth is not the same as Dropout with noise_shape=(1, 1, 1), as suggested in the tensorflow EfficientNet implementation (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/applications/efficientnet.py#L511). 
Using dropout, the branch is simply kept as is during inference time, however for stochastic depth, the branch is re-scaled based on its survival probability, before being merged with the main network.


Discussion:
Sounds good to me. Feel free to open an PR and request my review 😄 Also, welcome to ping me if you have any problem about testing suite and styles. Thank you.
@WindQAQ I have two implementations that work - one is a wrapper, that wraps the addition layer, before the residual branch is merged with the main network

The other is simply a layer that either replaces the Addition itself, or alternatively could just be placed last in the residual branch, scaling the branch before the normal add

They seem to be equally as performant, and both have the same regularizing effect (see this CIFAR10 ResNet32 benchmark https://github.com/MHStadler/tf_keras_playground/blob/master/Cifar10_StochDepth_Showcase.ipynb)

So I guess it just depends on which one fits better - I prefer the layer (since it's more straight forward, and also more flexible), but previous discussions were centered around the wrapper idea, so I figured I'd ask

The barebones layer implementation is available here: https://github.com/MHStadler/tf_keras_playground/blob/master/tf_keras_playground/layers/stochastic_depth.py
And the wrapper one here: https://github.com/MHStadler/tf_keras_playground/blob/master/tf_keras_playground/layers/wrappers.py

For how they are used in the model, see the model code; https://github.com/MHStadler/tf_keras_playground/blob/master/tf_keras_playground/models/resnet.py#L64
(0 means constant depth, 1 means wrapper, and 2 means layer stochastic depth)

Let me know which one you think would be better suited - note that both have the "issue" of still being updated with historical gradients, even if they are dropped. But my benchmark used NAG momentum like the original paper, and the regularizing effect remains, just without the performance gains of completely skipping the layers in the forward pass (looking at the original pytorch implementation, I don't think this is achieveable in tensorflow, without writing a custom training loop)
But this behavious is inline with how the very close related Dropout behaves in tensorflow, so I think it should be fine
> @WindQAQ I have two implementations that work - one is a wrapper, that wraps the addition layer, before the residual branch is merged with the main network
> 
> The other is simply a layer that either replaces the Addition itself, or alternatively could just be placed last in the residual branch, scaling the branch before the normal add
> 
> They seem to be equally as performant, and both have the same regularizing effect (see this CIFAR10 ResNet32 benchmark https://github.com/MHStadler/tf_keras_playground/blob/master/Cifar10_StochDepth_Showcase.ipynb)
> 
> So I guess it just depends on which one fits better - I prefer the layer (since it's more straight forward, and also more flexible), but previous discussions were centered around the wrapper idea, so I figured I'd ask
> 
> The barebones layer implementation is available here: https://github.com/MHStadler/tf_keras_playground/blob/master/tf_keras_playground/layers/stochastic_depth.py
> And the wrapper one here: https://github.com/MHStadler/tf_keras_playground/blob/master/tf_keras_playground/layers/wrappers.py
> 
> For how they are used in the model, see the model code; https://github.com/MHStadler/tf_keras_playground/blob/master/tf_keras_playground/models/resnet.py#L64
> (0 means constant depth, 1 means wrapper, and 2 means layer stochastic depth)
> 
> Let me know which one you think would be better suited - note that both have the "issue" of still being updated with historical gradients, even if they are dropped. But my benchmark used NAG momentum like the original paper, and the regularizing effect remains, just without the performance gains of completely skipping the layers in the forward pass (looking at the original pytorch implementation, I don't think this is achieveable in tensorflow, without writing a custom training loop)
> But this behavious is inline with how the very close related Dropout behaves in tensorflow, so I think it should be fine

Nice survey! Vote for `Layer` subclass. 
Okay - I guess Layer it is 

Shouldn't take too long to finish it up

Resolving patch:
diff --git a/tensorflow_addons/layers/__init__.py b/tensorflow_addons/layers/__init__.py
--- a/tensorflow_addons/layers/__init__.py
+++ b/tensorflow_addons/layers/__init__.py
@@ -38,4 +38,5 @@
 from tensorflow_addons.layers.tlu import TLU
 from tensorflow_addons.layers.wrappers import WeightNormalization
 from tensorflow_addons.layers.esn import ESN
+from tensorflow_addons.layers.stochastic_depth import StochasticDepth
 from tensorflow_addons.layers.noisy_dense import NoisyDense
diff --git a/tensorflow_addons/layers/stochastic_depth.py b/tensorflow_addons/layers/stochastic_depth.py
new file mode 100644
--- /dev/null
+++ b/tensorflow_addons/layers/stochastic_depth.py
@@ -0,0 +1,88 @@
+import tensorflow as tf
+from typeguard import typechecked
+
+
+@tf.keras.utils.register_keras_serializable(package="Addons")
+class StochasticDepth(tf.keras.layers.Layer):
+    """Stochastic Depth layer.
+
+    Implements Stochastic Depth as described in
+    [Deep Networks with Stochastic Depth](https://arxiv.org/abs/1603.09382), to randomly drop residual branches
+    in residual architectures.
+
+    Usage:
+    Residual architectures with fixed depth, use residual branches that are merged back into the main network
+    by adding the residual branch back to the input:
+
+    >>> input = np.ones((1, 3, 3, 1), dtype = np.float32)
+    >>> residual = tf.keras.layers.Conv2D(1, 1)(input)
+    >>> output = tf.keras.layers.Add()([input, residual])
+    >>> output.shape
+    TensorShape([1, 3, 3, 1])
+
+    StochasticDepth acts as a drop-in replacement for the addition:
+
+    >>> input = np.ones((1, 3, 3, 1), dtype = np.float32)
+    >>> residual = tf.keras.layers.Conv2D(1, 1)(input)
+    >>> output = tfa.layers.StochasticDepth()([input, residual])
+    >>> output.shape
+    TensorShape([1, 3, 3, 1])
+
+    At train time, StochasticDepth returns:
+
+    $$
+    x[0] + b_l * x[1],
+    $$
+
+    where $b_l$ is a random Bernoulli variable with probability $P(b_l = 1) = p_l$
+
+    At test time, StochasticDepth rescales the activations of the residual branch based on the survival probability ($p_l$):
+
+    $$
+    x[0] + p_l * x[1]
+    $$
+
+    Arguments:
+        survival_probability: float, the probability of the residual branch being kept.
+
+    Call Arguments:
+        inputs:  List of `[shortcut, residual]` where `shortcut`, and `residual` are tensors of equal shape.
+
+    Output shape:
+        Equal to the shape of inputs `shortcut`, and `residual`
+    """
+
+    @typechecked
+    def __init__(self, survival_probability: float = 0.5, **kwargs):
+        super().__init__(**kwargs)
+
+        self.survival_probability = survival_probability
+
+    def call(self, x, training=None):
+        if not isinstance(x, list) or len(x) != 2:
+            raise ValueError("input must be a list of length 2.")
+
+        shortcut, residual = x
+
+        # Random bernoulli variable indicating whether the branch should be kept or not or not
+        b_l = tf.keras.backend.random_bernoulli([], p=self.survival_probability)
+
+        def _call_train():
+            return shortcut + b_l * residual
+
+        def _call_test():
+            return shortcut + self.survival_probability * residual
+
+        return tf.keras.backend.in_train_phase(
+            _call_train, _call_test, training=training
+        )
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[0]
+
+    def get_config(self):
+        base_config = super().get_config()
+
+        config = {"survival_probability": self.survival_probability}
+
+        return {**base_config, **config}


Test patch:
diff --git a/tensorflow_addons/layers/tests/stochastic_depth_test.py b/tensorflow_addons/layers/tests/stochastic_depth_test.py
new file mode 100644
--- /dev/null
+++ b/tensorflow_addons/layers/tests/stochastic_depth_test.py
@@ -0,0 +1,58 @@
+import pytest
+import numpy as np
+import tensorflow as tf
+
+from tensorflow_addons.layers.stochastic_depth import StochasticDepth
+from tensorflow_addons.utils import test_utils
+
+_KEEP_SEED = 1111
+_DROP_SEED = 2222
+
+
+@pytest.mark.parametrize("seed", [_KEEP_SEED, _DROP_SEED])
+@pytest.mark.parametrize("training", [True, False])
+def stochastic_depth_test(seed, training):
+    np.random.seed(seed)
+    tf.random.set_seed(seed)
+
+    survival_probability = 0.5
+
+    shortcut = np.asarray([[0.2, 0.1, 0.4]]).astype(np.float32)
+    residual = np.asarray([[0.2, 0.4, 0.5]]).astype(np.float32)
+
+    if training:
+        if seed == _KEEP_SEED:
+            # shortcut + residual
+            expected_output = np.asarray([[0.4, 0.5, 0.9]]).astype(np.float32)
+        elif seed == _DROP_SEED:
+            # shortcut
+            expected_output = np.asarray([[0.2, 0.1, 0.4]]).astype(np.float32)
+    else:
+        # shortcut + p_l * residual
+        expected_output = np.asarray([[0.3, 0.3, 0.65]]).astype(np.float32)
+
+    test_utils.layer_test(
+        StochasticDepth,
+        kwargs={"survival_probability": survival_probability},
+        input_data=[shortcut, residual],
+        expected_output=expected_output,
+    )
+
+
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
+def test_with_mixed_precision_policy():
+    policy = tf.keras.mixed_precision.experimental.global_policy()
+
+    shortcut = np.asarray([[0.2, 0.1, 0.4]])
+    residual = np.asarray([[0.2, 0.4, 0.5]])
+
+    output = StochasticDepth()([shortcut, residual])
+
+    assert output.dtype == policy.compute_dtype
+
+
+def test_serialization():
+    stoch_depth = StochasticDepth(survival_probability=0.5)
+    serialized_stoch_depth = tf.keras.layers.serialize(stoch_depth)
+    new_layer = tf.keras.layers.deserialize(serialized_stoch_depth)
+    assert stoch_depth.get_config() == new_layer.get_config()