使用 tf.distribute.MirroredStrategy() 时无法加载 tensorflow keras 检查点

Question

我正在尝试从使用 ModelCheckpoint 回调创建的检查点加载 tf.keras (v1.15.0) 模型，通过删除多个层并添加新层来对其进行修改，然后继续在新任务上对其进行训练。我正在使用 tf.distribute.MirroredStrategy() 使用 2 个 gpus 进行分布式训练。

strategy = tensorflow.distribute.MirroredStrategy()
with strategy.scope():

    # Load pretrained model from checkpoint
    model = get_model()
    model.load_weights('file_name.hdf5')

    # Chop off some layers, add new layers
    model = modify_pretrained_model(model)

    model.compile(optimizer=opt, loss=loss)

模型可以正常加载并编译，我可以运行 model.summary()，但是当我调用 model.fit() 或 model.predict() 时，我的 python 堆栈中出现以下错误：

  (0) Failed precondition: Error while reading resource variable compression0_conv0_batchnorm/moving_variance from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/compression0_conv0_batchnorm/moving_variance/N10tensorflow3VarE does not exist.
     [[{{node time_distributed_1/model_1/compression0_conv0_batchnorm/FusedBatchNormV3/ReadVariableOp_1}}]]
     [[dense_1_1/Sigmoid/_225]]
  (1) Failed precondition: Error while reading resource variable compression0_conv0_batchnorm/moving_variance from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/compression0_conv0_batchnorm/moving_variance/N10tensorflow3VarE does not exist.
     [[{{node time_distributed_1/model_1/compression0_conv0_batchnorm/FusedBatchNormV3/ReadVariableOp_1}}]]
0 successful operations.
1 derived errors ignored

This issue 似乎解决了这个确切的问题，但没有使用 tf.distribute 继续训练。

当我在分发范围之外实例化会话并在分发范围内设置对它的引用时，代码会崩溃并出现相同的错误。

tf_config = some_custom_config
sess = tf.Session(config=tf_config)
graph = tf.get_default_graph()

strategy = tensorflow.distribute.MirroredStrategy()
with strategy.scope():

    set_session(sess)

    # Load pretrained model from checkpoint
    model = get_model()
    model.load_weights('file_name.hdf5')

    # Chop off some layers, add new layers
    model = modify_pretrained_model(model)

    model.compile(optimizer=opt, loss=loss)

Answer 1

我花了 2-3 天的时间试图弄清楚这一点。唯一真正有效的是升级到 tf 2.0.0。然后一切都像魔术一样工作。或者，作为最后的手段，我能够训练第一个模型，添加和删除额外的层，重新编译，并使用相同的分布策略在相同的 python 执行中继续训练，但永远无法使用分布策略重新加载 tf.keras ModelCheckpoint在 tf 1.15.0 中。