How to implement batch normalization layer for tensorflow multi-GPU code

Question

I have created a multi-GPU network as Cifar10_multigpu

In the inference implementation, they said:

We instantiate all variables using tf.get_variable() instead of tf.Variable() in order to share variables across multiple GPU training runs. If we only ran this model on a single GPU, we could simplify this function by replacing all instances of tf.get_variable() with tf.Variable().

So I did it about all my conv2d layers as the example, but what about batchnorm layer? How do I implement it by myself?

Can I use tensorflow.contrib.slim.batch_norm in this case? The example does not contain any recommendation about batch norm layer.

Answer 1

Simply use tf.layers.batch_normalization. It also creates variables via tf.get_variable(), hence they can be shared as well.

In addition, it works seamlessly with tf.layers.conv* functions.

Update: tf.nn.batch_normalization is fine too. It's a more low-level function that requires you manage mean and variance tensors yourself. In fact, tf.layers.batch_normalization is a wrapper over tf.nn.* functions, which also includes tf.nn.fused_batch_norm (a faster fused version).