#
27 Dec 2016
A Gentle Guide to Using Batch Normalization in Tensorflow

I recently made the switch to TensorFlow and am very happy with how easy it was to get things done using this awesome library. Tensorflow has come a long way since I first experimented with it in 2015, and I am happy to be back.

Since I am getting myself re-acquainted with TensorFlow, I decided that I should write a post about how to do batch normalization in TensorFlow. It’s kind of weird that batch normalization still presents such a challenge for new TensorFlow users, especially since TensorFlow comes with invaluable functions like `tf.nn.moments`

, `tf.nn.batch_normalization`

, and even `tf.contrib.layers.batch_norm`

. One would think that using batch normalization in TensorFlow will be a cinch. But alas, confusion still crops up from time to time, and the devil really lies in the details.

## Batch Normalization The Easy Way

Perhaps the easiest way to use batch normalization would be to simply use the `tf.contrib.layers.batch_norm`

layer. So let’s give that a go! Let’s get some imports and data loading out of the way first.

Next, we define our typical `fully-connected + batch normalization + nonlinearity`

set-up

One thing that might stand out is the `phase`

term. We are going to use as a placeholder for a boolean which we will insert into `feed_dict`

. It will serve as a binary indicator for whether we are in training `phase=True`

or testing `phase=False`

mode. Recall that batch normalization has distinct behaviors during training verus test time:

Training

Normalize layer activations according to mini-batch statistics.

During the training step, update population statistics approximation via moving average of mini-batch statistics.

Testing

Normalize layer activations according to estimated population statistics.

Do

notupdate population statistics according to mini-batch statistcs from test data.

Now we can define our very simple neural network for MNIST classification.

Now that we have defined our data and computational graph (a.k.a. model), we can train the model! Here is where we need to notice a very important note in the `tf.contrib.layers.batch_norm`

documentation:

Note: When is_training is True the moving_mean and moving_variance need to be updated, by default the update_ops are placed in tf.GraphKeys.UPDATE_OPS so they need to be added as a dependency to the train_op, example:

update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) if update_ops: updates = tf.group(*update_ops) total_loss = control_flow_ops.with_dependencies([updates], total_loss)

If you are comfortable with TensorFlow’s underlying graph/ops mechanism, the note is fairly straight-forward. If not, here’s a simple way to think of it: when you execute an operation (such as `train_step`

), only the subgraph components relevant to `train_step`

will be executed. Unfortunately, the `update_moving_averages`

operation is not a parent of `train_step`

in the computational graph, so we will never update the moving averages! To get around this, we have to explicitly tell the graph:

Hey graph, update the moving averages before you finish the training step!

Unfortunately, the instructions in the documentation are a little out of date. Furthermore, if you think about it a little more, you may conclude that attaching the `update`

ops to `total_loss`

may not be desirable if you wish to compute the `total_loss`

of the test set during test time. Personally, I think it makes more sense to attach the `update`

ops to the `train_step`

itself. So I modified the code a little and created the following training function

And we’re done! We can now train our model and see what happens. Below, I provide a comparison of the model without batch normalization, the model with pre-activation batch normalization, and the model with post-activation batch normalization.

As you can see, batch normalization really does help with training (not always, but it certainly did in this simple example).

## Additional Remarks

You have the choice of applying batch normalization either before or after the non-linearity, depending on your definition of the “activation distribution of interest” that you wish to normalize. It will probably end up being a hyperparameter that you’ll just have to tinker with.

There is also the question of what it means to *share* the same batch normalization layer across multiple models. I think this question should be treated with some care, and it really depends on what you think you’ll gain out of sharing the batch normalization layer. In particular, we should be careful about sharing the mini-batch statistics across multiple data streams if we expect the data streams to have distinct distributions, as is the case when using batch normalization in recurrent neural networks. As of the moment, the `tf.contrib.layers.batch_norm`

function does not allow this level of control.

But I do!… Kind of. It’s a fairly short piece of code, so it should be easy to modify it to fit your own purposes. **runs away**

End of post