11

我想知道在使用多 GPU 进行训练时，通过同步批量统计信息来实现批量标准化层的可能方法。

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

Torch我在这里和这里看到了一些评论，显示 running_mean 和 running_var 可以同步，但我认为 batch mean 和 batch var 不能或难以同步。

Tensorflow正常情况下和 caffe 和 torch 是一样的。BN的实现就是指这个。我知道 tensorflow 可以将操作分发到tf.device(). 但是means和vars的计算是在BN层的中间，所以如果我在cpu中收集means和vars，我的代码会是这样的：

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 可能是最好的选择，但我无法解决破解代码问题。与其他图书馆的解决方案也将受到欢迎。

4

3 回答 3

3

我不确定我是否完全理解您的问题，但如果您正确设置了变量范围，该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

你会看到一些慢下来（出于这个原因，我们通常只使用一个塔来计算批量规范统计），但它应该做你想要的。

于 2018-01-06T00:15:21.813 回答
2

自 TF2.2 以来，可以使用专门的 keras 层 SyncBatchNormalization https://www.tensorflow.org/api_docs/python/tf/keras/layers/experimental/SyncBatchNormalization

于 2020-09-01T17:06:27.977 回答
0

我想出了一种在纯 tensorflow 和纯 python 中实现同步批处理规范的方法。

该代码可以在 Cityscapes 上训练PSPNet并获得可比的性能。

于 2018-04-26T08:48:20.480 回答