Argument Scopes in Tensorflow

Intro

Tensorflow has a very useful feature called an argument scope. Argument scopes allow you to specify default values for tensorflow layer functions. This, in turn, allows you to write tensorflow code that is much simpler, more maintainable code that contains less clutter and repetition. This blog post will briefly cover how to make use of argument scopes.

Conv2d example

Creating neural network layers, especially when you specify your own non-defaut values, can take up a lot of space. For example, a convolutional layer might look something like this:

x = tf.contrib.layers.conv2d(X,
    num_outputs=32,
    kernel_size=3,
    stride=1,
    padding='VALID',
    activation_fn=tf.nn.elu,
    weights_initializer=tf.contrib.layers.xavier_initializer(),
    weights_regularizer=tf.contrib.layers.batch_norm)

Usually, you will want several such layers stacked one after another. Many of the arguments will also be exactly the same. This way of doing things takes up a lot of lines of code.

We can condense the code by creating argument scopes, that set default values for the conv2d function as follows:

with tf.contrib.framework.arg_scope(
    [tf.contrib.layers.conv2d],
    kernel_size=3,
    stride=1,
    padding='VALID',
    activation_fn=tf.nn.elu,
    weights_initializer=tf.contrib.layers.xavier_initializer(),
    weights_regularizer=tf.contrib.layers.batch_norm)
    ):
    x = tf.contrib.layers.conv2d(X, num_outputs=32)
    x = tf.contrib.layers.conv2d(x, num_outputs=32)
    x = tf.contrib.layers.conv2d(x, num_outputs=64, stride=2)

Now, each time we call conv2d, we only need to specify values for a small portion of the arguments that actually change between layer to layer.

We can also override the default values that were specified in the argument scope. In the last line of the code, we over-rode the stride argument to have a value of 2 instead of 1.

Default Arguments for Multiple Functions at once

You can also specify the default values of arguments for multiple functions at once, provided they all contain the argument names being specified in the argument scope.

For instance, in the following snippet of code, tf.contrib.layers.conv2d and tf.contrib.layers.fully_connected both share the arguments:

• activation_fn
• weights_initializer
• weights_regularizer

We can, therefore, specify defaults for those arguments for both functions as follows:

with tf.contrib.framework.arg_scope(
    [tf.contrib.layers.conv2d, tf.contrib.layers.fully_connected],
    activation_fn=tf.nn.elu,
    weights_initializer=tf.contrib.layers.xavier_initializer(),
    weights_regularizer=tf.contrib.layers.batch_norm
    ):
    x = tf.contrib.layers.conv2d(X, num_outputs=32,  kernel_size=3, stride=2)
    x = tf.contrib.layers.conv2d(x, num_outputs=32,  kernel_size=3, stride=2)
    x = tf.contrib.layers.flatten(x)
    x = tf.contrib.layers.fully_connected(x, num_outputs=32)
    Y = tf.contrib.layers.fully_connected(x, num_outputs=10, activation_fn=None)