将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 不同的结果？

我目前正在使用一些 BatchNorm 层训练一个基于 ResNet 的模型，如果我们只尝试在第一个等级上加载模型(并在其他等级上构建/编译模型(，我们会遇到一个停滞的张量问题 (https://github.com/horovod/horovod/issues/1271(。但是，如果我们在恢复时对所有等级调用hvd.load_model，训练开始正常恢复，但它似乎立即发散，所以我很困惑在所有等级上加载检查点模型(带有hvd.load_model(是否会以某种方式导致训练发散？但与此同时，由于 https://github.com/horovod/horovod/issues/1271，我们不能只在等级 0 上加载它，导致批次规范挂起在 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))
})