Question

我想了解在使用多GPU进行培训时，通过同步批量统计信息实现批量规范化层的可能方法。

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

Torch 我看过一些评论here和here，它们显示了running_mean和running_var可以同步，但我认为批量平均值和批量变量不能或者很难同步。

Tensorflow 通常情况下，它与caffe和torch相同。 BN的实施指的是this。我知道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)

这仅适用于一个BN图层。为了在cpu中收集统计信息，我必须打破代码。如果我有超过100个BN层，那将是麻烦的。

我不是那些图书馆的专家，所以也许会有一些误解，请随时指出我的错误。

我不太关心训练速度。我正在进行图像分割，其消耗大量GPU内存并且BN需要合理的批量大小（例如，大于16）以用于稳定的统计。因此使用多GPU是不可避免的。在我看来，tensorflow可能是最好的选择，但我无法解决破解代码问题。其他图书馆的解决方案也将受到欢迎。

Answer 1

我不确定我是否完全理解您的问题，但如果您正确设置了变量范围，tf.GraphKeys.UPDATE_OPS集合应自动为每个塔提供batch_norm的更新操作。如果同步应用所有update_ops，它们将由参数服务器隐式平均，您所要做的就是确保在平均和应用渐变之前应用更新。（如果我理解你的意图）。

由于范围可变，每组更新操作都会更新相同的变量，因此要同步更新操作，您需要做的就是在完整的更新操作集上控制梯度计算。您还应该将所有批量规范图层封装在一个name_scope中，以避免在UPDATE_OPS中抓取任何无关的操作。代码骨架如下：

update_ops = []
for i, device in enumerate(devices):
  with tf.variable_scope('foo', reuse=bool(i > 0)):
    with tf.name_scope('tower_%d' % i) as name_scope:
      with tf.device(device):
        # Put as many batch_norm layers as you want here
      update_ops.extend(tf.get_collection(tf.GraphKeys.UPDATE_OPS,
                                          name_scope))
# make gradient calculation ops here
with tf.device(averaging_device):
  with tf.control_dependencies(update_ops):
    # average and apply gradients.

如果您想在某些现有代码上尝试此操作，请尝试在此处删除if i == 0行：https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10_estimator/cifar10_main.py#L115

你会看到一些减速（我们通常只使用一个塔来计算批次规范统计数据），但它应该做你想要的。

Answer 2

我已经找到了在纯粹的tensorflow和纯python中实现sync batch norm的方法。

The code可以在Cityscapes上训练PSPNet并获得相应的表现。

Answer 3

自TF2.2起，可提供专用的keras层SyncBatchNormalization https://www.tensorflow.org/api_docs/python/tf/keras/layers/experimental/SyncBatchNormalization

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

3 个答案: