將 Horovod 與 tf.keras 一起使用時如何從檢查點恢復？

Question

注意：我使用的是 TF 2.1.0 和 tf.keras API。 我在 0.18 和 0.19.2 之間的所有 Horovod 版本中都遇到了以下問題。

從 tf.keras h5 檢查點恢復時，我們是否應該在所有等級上調用hvd.load_model() ，或者我們只應該在等級 0 上調用它並讓BroadcastGlobalVariablesCallback回調與其他工作人員共享這些權重？ 方法 1 是否不正確/無效，因為它會擾亂訓練或產生與方法 2 不同的結果？

I'm currently training a ResNet-based model with some BatchNorm layers, and if we only try to load the model on the first rank (and build/compile the model on the other ranks), we get a stalled tensor issue ( https: //github.com/horovod/horovod/issues/1271 ）。 但是，如果我們在恢復時在所有等級上調用hvd.load_model ，訓練開始正常恢復但它似乎立即發散，所以我很困惑是否在所有等級上加載檢查點 model （使用hvd.load_model ）會以某種方式導致訓練分歧？ 但同時，由於https://github.com/horovod/horovod/issues/1271 ，我們無法僅將其加載到 rank 0，導致 Batch Norm 在 horovod 中掛起。 使用 BatchNorm tf.keras 層時，是否有人能夠僅在等級 0 上成功調用hvd.load_model ？ 有人可以在這里提供一些提示嗎？

謝謝！

Answer 1

根據這個： https://github.com/horovod/horovod/issues/120 ，這是解決方案：

You should also be able to specify optimizer via custom object:
model = keras.models.load_model('file.h5', custom_objects={
    'Adam': lambda **kwargs: hvd.DistributedOptimizer(keras.optimizers.Adam(**kwargs))
})