tensorflow batch normalization after a restore - python-3.x

Question

Let us say that we create a small network:
tf.reset_default_graph()
layers = [5, 3, 1]
activations = [tf.tanh, tf.tanh, None]
inp = tf.placeholder(dtype=tf.float32, shape=(None, 2 ), name='inp')
out = tf.placeholder(dtype=tf.float32, shape=(None, 1 ), name='out')
isTraining = tf.placeholder(dtype=tf.bool, shape=(), name='isTraining')
N = inp * 1 # I am lazy
for i, (l, a) in enumerate(zip(layers, activations)):
N = tf.layers.dense(N, l, None)
#N = tf.layers.batch_normalization( N, training = isTraining) # comment this line
if a is not None:
N = a(N)
err = tf.reduce_mean((N - out)**2)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
opt = tf.train.AdamOptimizer(0.05).minimize(err)
# insert vectors from the batch normalization
tVars = tf.trainable_variables()
graph = tf.get_default_graph()
for v in graph.get_collection(tf.GraphKeys.GLOBAL_VARIABLES):
if all([
('batch_normalization' in v.name),
('optimizer' not in v.name),
v not in tVars ]):
tVars.append(v)
init = tf.global_variables_initializer()
saver = tf.train.Saver(var_list= tVars)
This is a simple NN generated for optimization. The only thing that I am currently interested in is batch optimization (the line that has been commented out). Now, we train this network, save it, restore its and calculate the error again, we do ok:
# Generate random data
N = 1000
X = np.random.rand(N, 2)
y = 2*X[:, 0] + 3*X[:, 1] + 3
y = y.reshape(-1, 1)
# Run the session and save it
with tf.Session() as sess:
sess.run(init)
print('During Training')
for i in range(3000):
_, errVal = sess.run([opt, err], feed_dict={inp:X, out:y, isTraining:True})
if i %500 == 0:
print(errVal)
shutil.rmtree('models1', ignore_errors=True)
os.makedirs('models1')
path = saver.save( sess, 'models1/model.ckpt' )
# restore the session
print('During testing')
with tf.Session() as sess:
saver.restore(sess, path)
errVal = sess.run(err, feed_dict={inp:X, out:y, isTraining:False})
print( errVal )
Here is the output:
During Training
24.4422
0.00330666
0.000314223
0.000106421
6.00441e-05
4.95262e-05
During testing
INFO:tensorflow:Restoring parameters from models1/model.ckpt
5.5899e-05
On the other hand, when we uncomment the batch normalization line, and redo the above calculation:
During Training
31.7372
1.92066e-05
3.87879e-06
2.55274e-06
1.25418e-06
1.43078e-06
During testing
INFO:tensorflow:Restoring parameters from models1/model.ckpt
0.041519
As you can see, the restored value is far from what the model is predicting. Is there anything that I am doing wrong?
Note: I know that for batch-normalization I need to generate mini batches. I have skipped all of that to keep the code simple and yet complete.

Answer 1

Batch normalization layer, as defined in Tensorflow, needs to have access to the placeholder isTraining (https://www.tensorflow.org/api_docs/python/tf/layers/batch_normalization). Make sure you include it when you define the layer: tf.layers.batch_normalization(..., training=isTraining, ...).
The reason for this is that Batch Normalization Layers have 2 trainable parameters (beta and gamma) that are trained normally with the rest of the network, but they also have 2 extra parameters (batch mean and variance) that require you to tell them to train. you do this simply by aplying the recipe above.
Right now your code seems not to be training mean and variance. Instead, they are randomly fixed and the network is optimized with those. Later on, when you save and restore, they are reinitialized with different values, hence the network doesn't perform as it used to.