每个 epoch 后的 TensorFlow 分布式训练暂停

Question

我正在使用 Tensorflow MirroredStrategy 在 2 个 GPU 上并行训练一个神经网络。使用单个 GPU，每个 epoch 需要 19 秒才能完成，而使用 2 个 GPU，每个 epoch 需要 13 秒才能完成。我对此并不感到惊讶，因为我知道由于在训练期间更新变量的 all_reduce 开销，缩放并不完美。

但是，在分布式训练的每个 epoch 之后，都会有大约 8 秒的停顿。使用单个 GPU 时，此暂停时间小于 1 秒。有谁知道为什么在分布训练时每个 epoch 后会有这么长时间的停顿？

另外，任何人都可以解释在 epoch 结束时分布式训练中发生的不同情况吗？

Answer 1

显然这与在图形模式下运行 TF 有关。通过设置tf.compat.v1.enable_eager_execution()，问题就消失了。这也修复了导致问题的内存泄漏，因此暂停可能是由于 TF 复制了我没想到的内容。