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

Question

我已经实现了一种分布式策略来在多个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告诉我这是不可能的，并且梯度必须应用于CCD_。这对我来说也是有道理的，但我想知道是否有可能积累梯度并使用镜像策略？

Answer 1

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

更多详细信息可在此处的文档中找到。