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

Question

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

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

手电筒我看到了一些评论 here和 here ，这表明 running_mean 和 running_var 可以同步，但我认为批处理均值和批处理 var 不能或难以同步。

Tensorflow 通常，它与caffe和torch相同。 BN的实现引用this .我知道 tensorflow 可以将操作分发到由 tf.device() 指定的任何设备.但是mean和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 可能是最好的选择，但我无法解决破解代码问题。其他图书馆的解决方案也将受到欢迎。

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

你会看到一些速度变慢(出于这个原因，我们通常只使用一个塔来计算批量规范统计数据)，但它应该可以做你想要的。