Pitfalls of Batch Norm in TensorFlow and Sanity Checks for Training Networks

Dishank Bansal — Tue, 26 Jun 2018 07:33:45 GMT

In this post I will write about mistakes we make when writing our deep learning model and some sanity checks, which will help you to narrow down the little daemons and save your nights’ sleep.

I. Let’s get done with Batch Norm in TensorFlow, once and for all

Brief Theory of Batch Normalization

This revolutionary technique is introduced by Sergey Ioffe, Christian Szegedy in the paper , and this is cited for 4994 times as of now. According to the paper, batch normalization reduces the internal covariance shift i.e. it makes the learning of layers in the network more independent of each other. The objective of batch norm layer is to make input to the activation layer, unit Gaussian, so that neuron does not get saturate in case of sigmoid and tanh. It helps in fast convergence of network, allows you to be care free about weight initialization, works as regularization in place of dropout and other regularization techniques. Now I won’t be discussing much of theoretical details but more on implementation. I will take you through implementation of batch norm in TensorFlow and its mathematics.
At Training Time -

At training time. Images taken from Batch Normalization paper.

Now suppose that batch norm layer is used between convolution layer and activation, then x is the output of convolution layer, and y is output of batch norm layer and input to activation layer, which may be ReLU, sigmoid, etc. Batch Norm layer learns the parameter γ and β, and y depends on their values. Hence, it is shown that best representation of x is learnt using γ and β. Please, note other two statistics which are E[x] and Var[x], population mean and variance. They are estimated through moving mean and moving variance during training.

At Inference Time -

At inference time

Forward pass through batch norm layer at inference is different than at training. At inference, instead of batch mean(μ) and variance(σ2) we use population mean(E[x]) and variance(Var[x]) to calculate x^. Suppose you give batch of size one during inference and normalize using batch mean and batch variance, in that case x^=0 as μ=x then y=β, which shows that layer output(y) will be same for any input(x) which is wrong. Hence to normalize at inference time we use population mean and variance which we were calculating during training using moving mean and variance. In the next section, we will go through the implementation of Batch Norm in TensorFlow.

Caveats of Batch Norm in TF

Moving mean and variance not updating:
TensorFlow is graph based, calculations is performed in order of dependencies.

import tensorflow as tf

a = tf.placeholder(tf.float32, shape=[1,3])
b = tf.placeholder(tf.float32, shape=[1,3])
c = tf.Variable(tf.zeros([1, 3]), dtype=tf.float32)


def train(a,b):
    c = tf.add(c,a)
    e = some_train_op(a,b)
    return e

output = train(a,b)
loss = tf.log(output)

sess = tf.Session()
saver = tf.train.Saver()
sess.run(tf.global_variables_initializer())

for itr in range(1000):
    ls = sess.run(loss, {a:x, b:y})

saver.save(sess, some_path)

#Note that this code is for reference only.

Now, in this we run the loss tensor in each iteration, for every iteration train function will be called. c is a variable to sum the value of a for all iteration. So, when we save the model, we expect variable c to be equal to summation of all values passed to placeholder a in each iteration.

Let’s analyze in detail, the dependencies in graph of this model. When we run loss tensor, output will be computed then output will call the train from which value of e will be needed which will be calculated using a and b. In all this computation c won’t be evaluated at all as value of loss tensor does not depend on c. Solution is to add c as dependency before computing e. Therefore, solution

def train(a,b):
    c = tf.add(c,a)
    with tf.control_dependencies([c]):
        e = some_train_op(a,b)
    return e

When training the network we call the optimizer op which does all the computations and updates all the weights, but optimizer op never depends upon computation of moving mean or variance. Hence, moving mean and variance never get updates and remain at their initialized value.

Solution: TensorFlow provide major three functions for batch normalization

tf.nn.batch_normalization() - Never use this for training, it just compute the y = gamma*x^hat + beta where x^hat = (x-mean)/std_dev.
It doesn't account for keeping moving mean and variance.

tf.layers.batch_normalization() - This function can be used for your model. Note: Always add batch norm dependency when using this, either on optimizer op or any other op which you sure that will execute when optimizer op is run. ex:-

update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)                            with tf.control_dependencies(update_ops):                                                            
    self.solver = tf.train.AdamOptimizer(learning_rate=1e-04) 
                  .minimize(self.loss)

Here tf.GraphKeys.UPDATE_OPS is the collections which have moving mean and variance op.

tf.contrib.layers.batch_norm() - This function is in contrib module of TensorFlow. This function gives two ways to add dependencies as follows:
1.It has 'update' parameter, set update=None while calling the function and dependency will be added inside the function itself and you don't have to anything else.
Note:This is believed to be slow in comparison to the next method.

2. Let the update parameter have it's default value which is "tf.GraphKeys.UPDATE_OPS". Now here also you have to manually add dependency as above with tf.layers.batch_normalization().

Hence, you can use either tf.layers.batch_normalization() or tf.contrib.layers.batch_norm() with proper update dependency handling.

Different usage at train and test time:
From train to test time, we have to switch between batch mean and variance or population mean and variance. We have parameters is_trainingand training in tf.contrib.layers.batch_norm() and tf.nn.batch_normalization() respectively. When its value is True, batch mean & variance is used and moving mean and variance is updated for estimating population mean and variance. On False, population mean and variance is used. Hence you should use a placeholder train_phase for this. So, solution looks like the following

def Bn(x, is_train):
        return tf.contrib.layers.batch_norm(x, decay= 0.99,                                                                                                               
          is_training=is_train,center= True, scale=True, 
          reuse= False)

train_phase = tf.placeholder(tf.bool, name="is_training")
data = tf.placeholder(tf.float32, shape=[None,240,240,3],                                                                                                        name="Image")

bn_layer = Bn(data, train_phase)

loss = some_loss_op()

solver = tf.train.AdamOptimizer().minimize(loss)

# Training iteration
sess.run(solver, {data: input_data, train_phase:True})

#Validation iteration
sess.run(loss, {data: input_data, train_phase:False})

Sharing Batch Norm parameters:
Often you share weights in two parallel networks such as siamese or in discriminators of GANs. When you share your weights, batch norm parameters also get shared. Sometimes, sharing batch norm parameters across two networks can not be healthy, I will tell you exactly when.
Let’s say we are training a GAN, it will have two discriminator instance with share parameters, when the training starts, distributions of fake and real images are different. Now here we are updating the moving mean and variance with two different distributions, one is of fake images and the other real one. But as the training continues and generator learns the distribution of real data, then distribution of both fake and real data eventually becomes approximately same. Then we are updating the mean and variance with same distribution, so our population mean and variance should be estimated statistics of real data after complete training. Hence, theoretically, sharing batch norm parameter may not be such a big problem in case of GANs. But I will do some experiments in both the conditions and update about the empirical results with new blog.
Now this blog post can be good explanation with better visualization for cases where sharing batch norm parameters can adversely affect the accuracy of the model.
In a nutshell, sharing batch norm parameters is okay until input distributions to both shared batch norm layers are same, if this is not the case try to use reuse=False in above mentioned batch norm functions, which will allow you to have independent batch norm parameters for both of your networks in which you wanted to share weights.

II. Sanity Checks and Debugging -

End to End Image Models:
In tasks like image segmentation or depth prediction that have image as an input, we expect a equally sized image as an output. Train your model to learn identity function i.e. give input image as the target image. In this case the network will represent a identity function. If this is trained perfectly but still you are unable to learn segmentation or depth, may be your model doesn’t have enough capacity. You can try to add some more layers or may be better loss function for your task.
Target and output range:
If you have normalized your target image in 0 to 1, then make sure in the image generating layer to have sigmoid as activation function and not tanh, as tanh will have range -1 to 1 which will make learning difficult. If normalized between -1 to 1, then use tanh.
Fluctuating Loss:
Most probable reason can be that your model doesn’t have enough capacity. Hence weights changes rapidly to learn the model but due to limited capacity they keep changing their weights rapidly but never getting to close to optimum. You should also try decreasing the learning rate.
Hyper-parameters:
Never ignore hyper-parameters. They may not be very crucial for accuracy and performance of your network in general, but are definitely very important for convergence and stability of training. Negligent weight initialization may cause exploiting/vanishing gradients or NaNs in your network. Check the std_dev of your weight initialization in case you are using normal distribution, range [0.2, 0.01] is okay. Learning rate and batch size are first thing to experiment with if loss is behaving abnormally.
Shuffling of sequential data:
When your input data is sequential which most often happen with videos such as in autonomous driving like segmentation, depth prediction, etc. You should consider shuffling the data, training with sequential data can adversely affect model accuracy at testing.
Note: You should have sequential data whenever task requires temporal consistency like optical flow, visual odometry, etc.
Scaling the losses:
In case you are using more than one type of loss in your network such as MSE, adversarial, L1, feature loss, SSIM then make sure all losses are scaled properly to be of same order i.e. of mse is in order of 1e-01 and feature loss is of order of 1e03, then scale the feature loss to be of same order.
Graph Visualization, Greatest weapon in this war:
If you re using TF and other framework having visualization tool. You should use visualization to check for connections, they will help you to find out missed connection, dimension error or weight sharing error.
Sharing Batch Norm parameter:
As mentioned above, if you are sharing batch norm parameters in two networks then make sure that input data distribution of both network is same.

This blog post also have exhaustive list of other reasons for failure of your network.

Thank you all for reading this blog, I hope that it will save your struggle with batch norm and training networks. I welcome all your comments and do let me know if want to add anything or feel any content to be misleading. You can reach me at dishank.bansal@iitkgp.ac.in for your queries.

References:

Pitfalls of Batch Norm in TensorFlow and Sanity Checks for Training Networks was originally published in TDS Archive on Medium, where people are continuing the conversation by highlighting and responding to this story.

Stories by Dishank Bansal on Medium