在 Tensorflow 2 中使用分布式策略累积梯度答案

【问题标题】：Accumulate gradients with distributed strategy in Tensorflow 2在 Tensorflow 2 中使用分布式策略累积梯度
【发布时间】：2022-01-13 06:42:57
【问题描述】：

我已经实施了一种分布式策略来在多个 GPU 上训练我的模型。

strategy = tf.distribute.MirroredStrategy(devices=devices[:FLAGS.n_gpus])
strategy.run(fn=self.train_step, args=(model, data))

我的模型现在变得更复杂、更大，我不得不减少批量大小以适应 GPU。梯度现在非常嘈杂，我想通过累积梯度再次增加批量大小。

现在我的问题是：使用镜像策略时这是否可能？我知道损失和梯度无论如何都是在副本中组合的，所以有没有办法在副本中对它们求和，例如在批次上运行的循环？我尝试了直截了当的方法，并返回了每个副本计算的梯度，以便在 strategy.run() 之外添加和应用它们，如下所示：

for b in batches:
    per_replica_gradients = strategy.run(fn=self.train_step, args=(model, data))
    total_gradient += per_replica_gradients
optimizer.apply_gradients(zip(total_gradient, model.trainable_variables)

但 Tensorflow 告诉我这是不可能的，必须使用 strategy.run() 应用渐变。这对我来说也很有意义，但我想知道是否有可能累积梯度并使用镜像策略？

【问题讨论】：

标签： python tensorflow tensorflow2.0 multi-gpu

【解决方案1】：

您可以使用tf.distribute.ReplicaContext.all_reduce：这与Strategy.reduce 的不同之处在于它用于副本上下文并且不会将结果复制到主机设备。 all_reduce 通常应用于训练步骤中的缩减，例如梯度。

更多细节可以在文档here中找到。

【讨论】：

哇，谢谢，我会试试的