tensorflow - 如何在 TensorFlow 的多 GPU 设置中使用批量归一化？

Question

在 How could I use Batch Normalization in TensorFlow? 上引用这篇文章.

我有一个类似于 CIFAR10 示例的多 GPU 设置。当我将 tf.contrib.layers.batch_norm 插入我的网络定义时，我在 average_gradients 中得到一个 NoneType 对象。具体来说，变量 g 是 NoneType。

def average_gradients(tower_grads):
    average_grads = []
    for grad_and_vars in zip(*tower_grads):
        grads = []
        for g, _ in grad_and_vars:
            expanded_g = tf.expand_dims(g, 0)
            grads.append(expanded_g)
        grad = tf.concat(0, grads)
        grad = tf.reduce_mean(grad, 0)
        v = grad_and_vars[0][1]
        grad_and_var = (grad, v)
        average_grads.append(grad_and_var)
    return average_grads

有关如何在多 GPU 环境中运行批归一化的一些示例代码会有所帮助。

编辑:

只需删除“batch_norm”变量即可解决此错误。然而，这里最紧迫的问题是，每个 Batch Normalization 在每个 GPU 上都有一个 beta 和 gamma，它们有自己的移动平均值。在推理时如何解决所有 GPU 上的所有这些移动平均数？

Answer 1

只需跨 GPU 独立使用 BN，同时使用其中一种塔式方法来更新移动均值。

with tf.device('..'):
  x,y = iterator.get_next()

  // NN with variables copied over to each of the GPUs
  loss = tower_loss(..)

  // use last tower statistics to update the moving mean/variance 
  batchnorm_updates = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope=scope)

apply_gradient_op = average_gradients(*grads)
batchnorm_updates_op = tf.group(*batchnorm_updates)
train_op = tf.group(apply_gradient_op, batchnorm_updates_op)

从多条评论中收集到 here ，这种简单的异步方法在实践中适用于大多数领域，但语义分割、 Action 视频识别等问题除外，在这些问题中，批量大小非常小，异步 BN 无法提供通常情况下的速度提升.