使用同步手段和变量来实现多GPU BN层的方法

我想知道在使用多GPU进行训练时,通过同步批量统计来实现批量规范化层的可能方法.

Caffe也许有一些caffe可以做的变种,比如链接.但是对于BN层,我的理解是它仍然只同步层的输出,而不是平均值和变量.也许MPI可以同步手段和变量,但我认为MPI有点难以实现.

火炬我在这里和这里看到了一些评论,它们显示了running_mean和running_var可以同步,但我认为批量平均值和批量变量不能或难以同步.

Tensorflow通常,它与caffe和torch相同.BN的实施是指这一点.我知道tensorflow可以将操作分配给指定的任何设备tf.device().但是平均值和变量的计算是在BN层的中间,所以如果我在cpu中收集平均值和变量,我的代码将是这样的:

cpu_gather = []
label_batches = []
for i in range(num_gpu):
    with tf.device('/gpu:%d' % i):
        with tf.variable_scope('block1', reuse=i > 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 > 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)