使用同步均值和方差的多GPU BN层实现方法

Question

我想知道在使用多个GPU进行训练时，实现批量归一化层并同步批量统计的可能方法。

Caffe 或许有一些可以实现的caffe变体，比如这里。但对于BN层，我的理解是它仍然只同步层的输出，而不是均值和方差。也许MPI可以同步均值和方差，但我认为MPI有点难以实现。

Torch 我看到一些评论这里和这里，显示running_mean和running_var可以同步，但我认为batch mean和batch 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)

这只是针对一个BN层的情况。为了在CPU中收集统计信息，我必须打破代码。如果我有超过100个BN层，那将会很麻烦。

我不是这些库的专家，所以可能存在一些误解，请随意指出我的错误。

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

- LI Xuhong

---

任何评论都受到赞赏。 - LI Xuhong

似乎在caffe中sync_bn_layer可以做到这一点。 - LI Xuhong

Answer 1

我不确定我是否完全理解了您的问题，但只要您正确设置变量范围，tf.GraphKeys.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。您会看到一些减速（出于这个原因，我们通常只使用一个塔来计算批量归一化统计信息），但它应该能够实现您想要的功能。

- Eli Bixby

---

谢谢Eli Bixby。您是否成功地使用多个GPU训练了BN？请查看我的问题并给我一些评论：https://stackoverflow.com/questions/48150720/how-to-update-variable-of-batchnorm-in-multiple-gpus-in-tensorflow - Jame

感谢@Eli Bixby的回答，但很抱歉可能有一些错误。对于BN，不仅应该累积或更新“statistics”，每次反向传播的梯度也应该被考虑在内。如果使用小批量大小，则梯度根本不稳定。您在此提出的只是前向传递，反向传递是通过自动微分计算的，并且没有正确执行。 - LI Xuhong

你应该使用足够大的批量大小，以使每个GPU都饱和。根据我们的经验，即使在分片GPU之间进行分片，这个批量大小也足够大，可以通过仅从单个GPU分片传播梯度来产生稳定的batch_norm梯度（请参见我链接的代码）。据我所知，目前高级批量规范函数中没有一种方法可以在应用之前平均批量规范梯度，您必须在低级TF中自己实现它。 - Eli Bixby

@Eli Bixby 谢谢。我处理分割任务，通常会消耗大量内存，因此无法使用大批量大小。您是否有在低级TF中实现的经验或一些具体示例？最好有计算梯度的实现。 - LI Xuhong

我没有尝试手动完成这个特定任务的经验，但它应该遵循与手动平均梯度完全相同的模式，只是您需要平均批量归一化统计数据，以计算梯度。因此，您需要手动创建一些批量归一化操作并将其固定到设备上，然后使用操作列表提取平均统计信息。这类似于我上面链接的代码（用于梯度）。抱歉我不能更有帮助！如果我有时间，我可能会尝试自己做这个。 - Eli Bixby

Answer 2

自 TensorFlow 2.2 版本起，我们提供了一种专门的 Keras 层——SyncBatchNormalization。您可以通过以下链接了解更多信息：https://www.tensorflow.org/api_docs/python/tf/keras/layers/experimental/SyncBatchNormalization。

- Henrique Mendonça

Answer 3

我已经找到了一种在纯TensorFlow和纯Python中实现同步批量归一化的方法。

这段代码使得在Cityscapes上训练PSPNet并获得可比较的性能成为可能。

- LI Xuhong