在tf.keras中使用Horovod时，如何从检查点恢复？

Question

注意:我使用的是TF 2.1.0和tf.keras接口。我在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))
})