Tensorflow 每個 epoch 后的分布式訓練暫停

Question

我正在使用 Tensorflow MirroredStrategy 在 2 個 GPU 上並行訓練神經網絡。 使用單個 GPU，每個 epoch 需要 19 秒才能完成，而使用 2 個 GPU，每個 epoch 需要 13 秒才能完成。 我對此並不感到驚訝，因為我知道由於在訓練期間更新變量的 all_reduce 開銷，縮放並不完美。

但是，在分布式訓練的每個 epoch 之后，都會有大約 8 秒的停頓。 使用單個 GPU 時，此暫停時間小於 1 秒。 有誰知道為什么在分布訓練時每個 epoch 后會有這么長時間的停頓？

或者，任何人都可以解釋在一個時代結束時分布式訓練中發生的不同情況嗎？

Answer 1

顯然這與在圖形模式下運行 TF 有關。 通過設置tf.compat.v1.enable_eager_execution()問題消失了。 這也修復了導致問題的 memory 泄漏，因此暫停可能是由於 TF 復制了我沒想到的東西造成的。