如何在TensorFlow中的多gpu设置中使用批量标准化?

时间:2017-01-24 02:33:24

标签: tensorflow

How could I use Batch Normalization in TensorFlow?上引用此帖子。

我有一个类似于CIFAR10示例的multi-gpu设置。当我将tf.contrib.layers.batch_norm插入到我的网络定义中时,我在average_gradients中获得了一个NoneType对象。具体来说,variable 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

有关如何在multi-gpu环境中运行Batch Normalization的一些示例代码会有所帮助。

编辑:

只需删除“batch_norm”变量就可以解决这个问题。然而,这里的紧迫问题是每个批量标准化在每个GPU上都有一个beta和gamma,具有自己的移动平均值。所有这些移动平均线如何通过推理得到解决?

1 个答案:

答案 0 :(得分:0)

只需在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中收集到的结果来看,这种简单的异步方法在大多数域中都可以很好地实践,除了诸如语义分段,动作视频识别等问题(批处理规模极小且BN异步)无法承受正常的速度提升。