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

Question

I'm trying to load a tf.keras (v1.15.0) model from a checkpoint created with the ModelCheckpoint callback, modify it by removing several layers and adding new ones, and then continue training it on a new task. I'm using the tf.distribute.MirroredStrategy() to do distributed training with 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)

The model loads fine and compiles and I can run model.summary(), but when I call model.fit(), or model.predict() I get the following errors in my python stack:

  (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 seems to fix this exact issue but without using tf.distribute to continue training.

When I instantiate a session outside of the distribute scope, and set a reference to it inside the distribute scope the code crashes with the same error.

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

I spent a good 2-3 days trying to figure this out. The only thing that really worked was upgrading to tf 2.0.0. Then everything worked like magic. Alternatively as a last resort, I was able train the first model, add and remove additional layers, recompile, and continue training in the same python execution with the same distribution strategy, but was never able to reload a tf.keras ModelCheckpoint using distribution strategies in tf 1.15.0.