Grover模型的多GPU训练

Question

我正在尝试使用tensorflow == 1.13.1在多个GPU上训练Grover的鉴别器模型。该模型的默认配置假定对多个TPU进行训练，并被证明可以在具有多个TPU的系统上工作。但是，由于TPU不可用，我试图通过使用Tensorflow分配策略更改模型的运行配置，从而在多个GPU上分配训练，如下所示：

dist_strategy = tf.contrib.distribute.MirroredStrategy(
        num_gpus=8,
        cross_device_ops=AllReduceCrossDeviceOps('nccl', num_packs=8),
        # cross_device_ops=AllReduceCrossDeviceOps('hierarchical_copy')
)

然后我尝试如下更改GPU的运行配置：

 run_config = RunConfig(
        train_distribute=dist_strategy,
        eval_distribute=dist_strategy,
        log_step_count_steps=log_every_n_steps,
        model_dir=FLAGS.output_dir,
        save_checkpoints_steps=FLAGS.iterations_per_loop)

它似乎可以正确分配数据（我可以通过GPU使用情况和其他指标来判断），但是在收集渐变时它会失败，并出现以下错误：

ValueError: You must specify an aggregation method to update a MirroredVariable in Replica Context.

但是我在分发策略中提供了跨设备操作。在我看来，最终的解决方案是使用分配策略代码修改损失函数，但我认为它应该可以与TPU一起使用。有帮助吗？