<p>I've figured out a way to implement <a href="https://github.com/holyseven/PSPNet-TF-Reproduce/blob/master/model/utils_mg.py#L177" rel="nofollow noreferrer">sync batch norm</a> in pure tensorflow and pure python. </p>

<p><a href="https://github.com/holyseven/PSPNet-TF-Reproduce" rel="nofollow noreferrer">The code</a> makes it possible to train <a href="https://arxiv.org/abs/1612.01105" rel="nofollow noreferrer">PSPNet</a> on Cityscapes and get comparable performance.</p>
Ways to implement multi-GPU BN layers with synchronizing means and vars
<p>I'd like to know the possible ways to implement batch normalization layers with synchronizing batch statistics when training with multi-GPU. </p>

<p><strong>Caffe</strong> Maybe there are some variants of caffe that could do, like <a href="https://github.com/yjxiong/caffe" rel="noreferrer">link</a>. But for BN layer, my understanding is that it still synchronizes only the outputs of layers, not the means and vars. Maybe MPI can synchronizes means and vars but I think MPI is a little difficult to implemnt.</p>

<p><strong>Torch</strong> I've seen some comments <a href="https://github.com/tensorflow/tensorflow/issues/7439" rel="noreferrer">here</a> and <a href="https://github.com/torch/nn/issues/1071" rel="noreferrer">here</a>, which show the running_mean and running_var can be synchronized but I think batch mean and batch var can not or are difficult to synchronize.</p>

<p><strong>Tensorflow</strong> Normally, it is the same as caffe and torch. The implementation of BN refers <a href="https://github.com/ppwwyyxx/tensorpack/blob/master/tensorpack/models/batch_norm.py" rel="noreferrer">this</a>. I know tensorflow can distribute an operation to any device specified by <code>tf.device()</code>. But the computation of means and vars is in the middle of BN layer, so if I gather the means and vars in cpu, my code will be like this:</p>



<pre class="lang-py prettyprint-override"><code>cpu_gather = []
label_batches = []
for i in range(num_gpu):
    with tf.device('/gpu:%d' % i):
        with tf.variable_scope('block1', reuse=i &gt; 0):
            image_batch, label_batch = cifar_input.build_input('cifar10', train_data_path, batch_size, 'train')
            label_batches.append(label_batch)

            x = _conv('weights', image_batch, 3, 3, 16, _stride_arr(1))
            block1_gather.append(x)

with tf.device('/cpu:0'):
    print block1_gather[0].get_shape()
    x1 = tf.concat(block1_gather, 0)
    # print x1.get_shape()
    mean, variance = tf.nn.moments(x1, [0, 1, 2], name='moments')

for i in range(num_gpu):
    with tf.device('/gpu:%d' % i):
        with tf.variable_scope('block2', reuse=i &gt; 0):
            shape = cpu_gather[i].get_shape().as_list()
            assert len(shape) in [2, 4]
            n_out = shape[-1]
            beta, gamma, moving_mean, moving_var = get_bn_variables(n_out, True, True)

            x = tf.nn.batch_normalization(
                cpu_gather[i], mean, variance, beta, gamma, 0.00001)

            x = _relu(x)
</code></pre>

<p>That is just for one BN layer. For gathering statistics in cpu, I have to break the code. If I have more than 100 BN layers, that will be cumbersome. </p>

<p>I am not expert in those libraries so maybe there are some misunderstanding, feel free to point out my errors. </p>

<p>I do not care much about training speed. I am doing image segmentation which consumes much GPU memory and BN needs a reasonable batch size (e.g. larger than 16) for stable statistics. So using multi-GPU is inevitable. In my opinion, tensorflow might be the best choice but I can't resolve the breaking code problem. Solution with other libraries will be welcome too.</p>