Feature Request : Stochastic Depth

[alexanderalang]

Stochastic Depth (aka layer dropout) has been shown to speed up and improve training in ResNets, as well as overall accuracy on testing sets. Essentially, every training step a random subset of residual layers are entirely removed from the network, and training proceeds on the remaining layers. Direct connections are made between the missing layers.

It is described in this paper: https://arxiv.org/pdf/1603.09382.pdf. (Deep Networks with Stochastic Depth by Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, Kilian Q. Weinberger)

I can't think of a way to implement this with the python API without reconstructing the model every training iteration, and I'm not familiar the with the C++ API / Cudnn to try to write the op myself.

Of course I'm willing to try any python-only suggestions.

Thanks in advance,

Alex

[asimshankar]

We currently don't have anybody working on this. It would be great if you could help us by working on this and submitting a PR. Let us know if you need further clarification. Thanks!

[yaroslavvb]

It's possibly to do in pure Python. You can "skip" layers in forward prop by using tf.cond to decide which layer is used as an input to given layer.

You can do even more advanced things without leaving Python-land, ie, you can forward prop through some layers, and backprop through a different set of layers.
To skip layers for backprop, the basic trick is to use stop_gradient to control which implementation is used for backprop run. Here's another example

[alexanderalang]

@yaroslavvb thanks for the tip. Can I post my code back in this thread when I come up with something?

[yaroslavvb]

@alexanderalang yes, I think "contributions welcome" means that people are welcome to link to solutions here

[alexanderalang]

this isn't very elegant, but it works.

def layer_wrapper(layer,layer_args,layer_kwargs={},switch=1):
    output = layer(*layer_args,**layer_kwargs) #normal output of layer
        
    forward = tf.cond(tf.cast(switch,tf.bool), #if layer is disabled, inputs are passed on to next layer. 
                     lambda: output,
                     lambda: layer_args[0]) # requires 1st value of args to be the input
        
    backward = tf.cond(tf.cast(switch,tf.bool), #gradient is stopped for backprop if layer is disabled
                     lambda: forward,
                     lambda: tf.zeros_like(forward))
        
    return backward + tf.stop_gradient(forward-backward)


def assign_switches(vec):
    with tf.variable_scope('layer_dropout', reuse=True):
        switches = tf.get_collection('switches')
    assert len(vec) == len(switches)
    return [switches[i].assign(vec[i]) for i in range(len(switches))]

def make_dropout_vars():
    with tf.variable_scope('layer_dropout', reuse=True):
        switches = tf.get_collection('switches')
        switch = tf.Variable(1,name="switches/"+str(len(switches)))
        drop_prob = tf.Variable(0.5,name="probability/"+str(len(switches)))
        tf.add_to_collection('switches',switch)
        tf.add_to_collection('dropout_probabilities',drop_prob)
        return switch

To mirror @yaroslavvb 's example

def add(x,name):
    return x+1
x = tf.ones((), name="x")
y0 = layer_wrapper(add,(x,"y0"),switch=make_dropout_vars())
y1 = layer_wrapper(add,(x,"y1"),switch=make_dropout_vars())

y = tf.stack([y0, y1], name="y")

z = tf.reduce_sum(y)

grad = tf.gradients(z, [x])[0]

sess = tf.Session()
sess.run(tf.global_variables_initializer())

print("\nValue %.1f, gradient %.1f"%tuple(sess.run([z, grad])))

sess.run(assign_switches([0,0]))
print("Value %.1f, gradient %.1f"%tuple(sess.run([z, grad])))

sess.run(assign_switches([0,1]))
print("Value %.1f, gradient %.1f"%tuple(sess.run([z, grad])))

It produces the same gradients as in the example, but the value depends on whether the add nodes are activated (since add adds 1 and not 0 like in the example).

so the output is

Value 4.0, gradient 2.0
Value 2.0, gradient 0.0
Value 3.0, gradient 1.0

[alexanderalang]

Ok, this fixes the problem of the gradient not being passed on to the previous layer during backprop if a layer is disabled.

import tensorflow as tf
from tensorflow.python.framework import ops

@tf.RegisterGradient("PassThroughGrad")
def _pass_through_grad(unused_op, grad):
    return grad, grad

def pass_through(inputs,name):
    with tf.Graph().as_default() as g:
        with g.gradient_override_map({name:"PassThroughGrad"}):
            I = tf.identity(inputs,name)
            return I

def layer_wrapper(layer,layer_args,layer_kwargs={},switch=1,id_op_num=None):
    id_op_name= 'my_id_op/{}'.format(id_op_num)
    output = layer(*layer_args,**layer_kwargs) #normal output of layer
        
    forward = tf.cond(tf.cast(switch,tf.bool), #if layer is disabled, inputs are passed on to next layer. 
                     lambda: output,
                     lambda: layer_args[0])
        
    backward = tf.cond(tf.cast(switch,tf.bool), #gradient is stopped for backprop if layer is disabled
                     lambda: forward,
                     lambda: pass_through(forward,id_op_name))
        
    return backward + tf.stop_gradient(forward-backward)


def dropout(prob):
    def dec(layer):
        def wrapperwrapper(*layer_args,**layer_kwargs):
            switch = make_dropout_vars(prob)
            return layer_wrapper(layer,layer_args,layer_kwargs,switch=switch,id_op_num=switch.name[-1])
        return wrapperwrapper
    return dec



def assign_switches(vec):
    with tf.variable_scope('layer_dropout', reuse=True):
        switches = tf.get_collection('switches')
    assert len(vec) == len(switches)
    return [switches[i].assign(vec[i]) for i in range(len(switches))]

def make_dropout_vars(prob=0.5):
    with tf.variable_scope('layer_dropout', reuse=True):
        switches = tf.get_collection('switches')
        switch = tf.Variable(1,name="switches/"+str(len(switches)))
        drop_prob = tf.Variable(prob,name="probability/"+str(len(switches)))
        tf.add_to_collection('switches',switch)
        tf.add_to_collection('dropout_probabilities',drop_prob)
        return switch

Then you just decorate your function with @dropout(0.5) or whatever probability you want.
At this point I'm just wondering if there is a safer way of storing/assigning switching and dropout_probability variables. I'm still a tensorflow noob so variable namespacing and reuse kind of scares me. Also i'll take any tips for the naming of the Identity ops in pass_through to make sure that they are all unique, the numbering method I'm currently using seems sketchy.

[samjabrahams]

Hi guys,

I wanted to chip in a bit. I may be completely missing something, but I believe tf.cond already sets the gradient for the unused path to zero. This should prevent unwanted updates to Variables. Here is some code from a presentation/demo I'm giving tomorrow:

Stochastic depth layer (with 1x1 convolution depth mapping, if needed). Not totally full-featured, but works for a simple illustration.

def stochastic_depth_conv2d(inputs, filters, kernel_size, keep_prob, padding='same', activation=tf.nn.relu, name=None):
    default_name = tf.get_default_graph().unique_name('stochastic_depth_conv')
    with tf.variable_scope(name, default_name):
        def full_layer():
            return tf.layers.conv2d(inputs, filters, kernel_size, padding=padding, activation=activation, name='conv')
        def skip_layer():
            if inputs.get_shape().as_list()[-1] != filters:
                return tf.layers.conv2d(inputs, filters, [1, 1], padding=padding, activation=activation, name='skip')
            else:
                return inputs
        pred = tf.random_uniform([]) < keep_prob
        return tf.cond(pred, full_layer, skip_layer), pred

Some code to print out gradients based on randomly switching some layers "on" or "off":

tf.reset_default_graph()
inputs = tf.placeholder(tf.float32, [None, 3, 3, 1], name='inputs')
keep_prob = tf.placeholder(tf.float32, [], name='keep_prob')
conv1, pred1 = stochastic_depth_conv2d(inputs, 3, [3, 3], keep_prob, name='conv1')
conv2, pred2 = stochastic_depth_conv2d(conv1, 3, [3, 3], keep_prob, name='conv2')
conv3, pred3 = stochastic_depth_conv2d(conv2, 3, [3, 3], keep_prob, name='conv3')
init = tf.global_variables_initializer()

opt = tf.train.GradientDescentOptimizer(0.05)
var_list = tf.trainable_variables()
grads = opt.compute_gradients(conv3, var_list)

with tf.Session() as sess:
    sess.run(init)
    feed_dict = {
        inputs: np.random.normal(size=[1, 3, 3, 1]),
        keep_prob: 0.5
    }
    g, p1, p2, p3 = sess.run([grads, pred1, pred2, pred3], feed_dict)
    print(p1, p2, p3)
    for var, grad_value in zip(var_list, g):
        grad, value = grad_value
        print('',var.op.name, grad.squeeze(), sep='\n')

Sample output from the above:

False True False

conv1/conv/kernel
[[[ 0.  0.  0.]
  [ 0.  0.  0.]
  [ 0.  0.  0.]]

 [[ 0.  0.  0.]
  [ 0.  0.  0.]
  [ 0.  0.  0.]]

 [[ 0.  0.  0.]
  [ 0.  0.  0.]
  [ 0.  0.  0.]]]

conv1/conv/bias
[ 0.  0.  0.]

conv1/skip/kernel
[ 0.88905168  1.15772355  3.13942909]

conv1/skip/bias
[-0.24839747  0.25073451  2.12950516]

conv2/conv/kernel
[[[[ 1.11243725  0.06358987  0.06358987]
   [ 1.97543406  0.11292107  0.11292107]
   [ 1.45493352  0.08316787  0.08316787]]

  [[ 0.06358987  0.          1.04884744]
   [ 0.11292107  0.          1.86251295]
   [ 0.08316787  0.          1.37176561]]

  [[ 0.          0.          0.06358987]
   [ 0.          0.          0.11292107]
   [ 0.          0.          0.08316787]]]


 [[[ 0.          0.          1.93598056]
   [ 0.          0.          3.43785858]
   [ 0.          0.          2.53202868]]

  [[ 1.87239075  1.87239075  0.06358987]
   [ 3.32493734  3.32493734  0.11292107]
   [ 2.44886065  2.44886065  0.08316787]]

  [[ 0.06358987  0.06358987  0.        ]
   [ 0.11292107  0.11292107  0.        ]
   [ 0.08316787  0.08316787  0.        ]]]


 [[[ 0.82354331  0.82354331  0.82354331]
   [ 1.4624244   1.4624244   1.4624244 ]
   [ 1.07709503  1.07709503  1.07709503]]

  [[ 0.82354331  0.          0.        ]
   [ 1.4624244   0.          0.        ]
   [ 1.07709503  0.          0.        ]]

  [[ 0.          0.          0.82354331]
   [ 0.          0.          1.4624244 ]
   [ 0.          0.          1.07709503]]]]

conv2/conv/bias
[ 4.  3.  5.]

conv3/conv/kernel
[[[[ 0.  0.  0.]
   [ 0.  0.  0.]
   [ 0.  0.  0.]]

  [[ 0.  0.  0.]
   [ 0.  0.  0.]
   [ 0.  0.  0.]]

  [[ 0.  0.  0.]
   [ 0.  0.  0.]
   [ 0.  0.  0.]]]


 [[[ 0.  0.  0.]
   [ 0.  0.  0.]
   [ 0.  0.  0.]]

  [[ 0.  0.  0.]
   [ 0.  0.  0.]
   [ 0.  0.  0.]]

  [[ 0.  0.  0.]
   [ 0.  0.  0.]
   [ 0.  0.  0.]]]


 [[[ 0.  0.  0.]
   [ 0.  0.  0.]
   [ 0.  0.  0.]]

  [[ 0.  0.  0.]
   [ 0.  0.  0.]
   [ 0.  0.  0.]]

  [[ 0.  0.  0.]
   [ 0.  0.  0.]
   [ 0.  0.  0.]]]]

conv3/conv/bias
[ 0.  0.  0.]

[alexanderalang]

@samjabrahams so theres no need for the gradient_override_map with the Identity op? Does your code actually skip the layer during backprop, or does the weight update get carried out all the same, just with gradients that are all zero? I guess what I'm trying to understand is, for a very deep network, will this method still save time / memory?

[samjabrahams]

@alexanderalang the initial tests look promising. Here's some sample code and respective timing outputs with a very simple CNN (no pooling, batch norm, etc) doing a single training step on dummy data/labels and batch size 32. This is from running on a 2013 macbook pro CPU, so focus on the relative times instead of the total time.

Stochastic depth

Training code

# Stochastic version
tf.reset_default_graph()
print('Building graph...')
inputs = tf.placeholder(tf.float32, [None, 228, 228, 3], name='inputs')
labels = tf.placeholder(tf.int32, None)
labels_onehot = tf.one_hot(labels, 10)
keep_prob = tf.placeholder(tf.float32, [], name='keep_prob')
conv, _ = stochastic_depth_conv2d(inputs, 32, [3, 3], keep_prob)
for i in range(10):
    conv, _ = stochastic_depth_conv2d(conv, 32, [3, 3], keep_prob)
flat = tf.reshape(conv, [-1, 32 * 228 * 228])
logits = tf.layers.dense(flat, 10)
loss = tf.nn.softmax_cross_entropy_with_logits(labels=labels_onehot, logits=logits)

train = tf.train.GradientDescentOptimizer(0.05).minimize(loss)

init = tf.global_variables_initializer()

print('Done.')

with tf.Session() as sess:
    sess.run(init)
    feed_dict = {
        inputs: np.random.normal(size=[32, 228, 228, 3]),
        # Adjust this parameter to see run speed adjust
        keep_prob: 0.5,
        labels: np.random.randint(0, 10, size=[32])
    }
    print('Timing stochastic depth run...')
    start_t = time.time()
    _ = sess.run(train, feed_dict)
    end_t = time.time()
    print('Done.')
    print(end_t - start_t)

Output:

Building graph...
Done.
Timing stochastic depth run...
Done.
13.219794034957886

"Normal" Conv2D layer

Training code:

# Standard CNN version
tf.reset_default_graph()
print('Building graph...')
inputs = tf.placeholder(tf.float32, [None, 228, 228, 3], name='inputs')
labels = tf.placeholder(tf.int32, None)
labels_onehot = tf.one_hot(labels, 10)
conv = tf.layers.conv2d(inputs, 32, [3, 3], padding='same', activation=tf.nn.relu)
for i in range(10):
    conv = tf.layers.conv2d(conv, 32, [3, 3], padding='same', activation=tf.nn.relu)
flat = tf.reshape(conv, [-1, 32 * 228 * 228])
logits = tf.layers.dense(flat, 10)
loss = tf.nn.softmax_cross_entropy_with_logits(labels=labels_onehot, logits=logits)

train = tf.train.GradientDescentOptimizer(0.05).minimize(loss)

init = tf.global_variables_initializer()

print('Done.')

with tf.Session() as sess:
    sess.run(init)
    feed_dict = {
        inputs: np.random.normal(size=[32, 228, 228, 3]),
        labels: np.random.randint(0, 10, size=[32])
    }
    print('Timing standard run...')
    start_t = time.time()
    _ = sess.run(train, feed_dict)
    end_t = time.time()
    print('Done.')
    print(end_t - start_t)

Output:

Building graph...
Done.
Timing standard run...
Done.
25.8157160282135

Adjusting the keep_prob feed value in the first code alters the average relative runtime as you'd expect.

keep_prob = 0.2 -->  5.950854063034058 seconds
keep_prob = 0.8 --> 18.41518807411194 seconds

[awav]

tensorflow/tensorflow#8817, @samjabrahams, I ended up with the same solution. But, it would be easier if trainable parameter of the get_variable was tensor typed.

[samjabrahams]

@awav The trainable parameter just puts that Variable object into the TRAINABLE_VARIABLES collection. I don't think TensorFlow is designed to have the var_list passed into Optimizer.compute_gradients() changed (the compute_gradients operation has the variables it updates statically defined).

Since we're using tf.cond already in order to implement layer skipping in the forward pass, I'm not sure that it's easier to explicitly block the gradients as opposed to taking advantage of it via tf.cond for free :)

[dynamicwebpaige]

Though this is outside of the scope of layers in TensorFlow core, I'm sure that Stochastic Depth would be a welcome addition to TensorFlow Addons.

Transitioning the issue to tensorflow/addons now; @seanpmorgan for visibility.

[seanpmorgan]

Thanks for transferring! I believe this could be implemented as a Keras Layer Wrapper should there still be a community need for this.