Trying to create optimizer slot variable under the scope for tf.distribute.Strategy which is different from the scope used for the original variable

Question

I'm trying to train, save and load a tensorflow model. An outline of my code is as follows:

devices = ["/device:GPU:{}".format(i) for i in range(num_gpus)]
strategy = tf.distribute.MirroredStrategy(devices)
with strategy.scope():
    # Create model
    model = my_model(some parameters)
    model.build((parameters))
    model.summary()

    # Adam optimizer with EMA smoothing
    opt = tf.keras.optimizers.Adam(learning_rate)
    opt = tfa.optimizers.MovingAverage(opt, average_decay=ema_decay)

    model.compile(
        optimizer=opt,
        loss=loss_dict,
        loss_weights=loss_weights,
        metrics=metrics_dict)


    #adding softmax layer
    output_list = list()
    for i in range(num_class):
        output_list.append(tf.keras.layers.Softmax(name=f"name_{str(i)}")(model.output[i]))
    output_list.append(model.output[num_class])

    
    model_b = tf.keras.Model(inputs=model.input, outputs=output_list)
    model_b.build((None, None, feats_dim, 1))
    model_b.compile(optimizer=opt)

model.fit(parameters... callbacks=[cp_callback, logger_callback, tb_callback])
model_b.load_weights(checkpoint_path)
model_b.save(os.path.join(model_path, "model.h5"))

The checkpoints are saved using:

cp_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_path, verbose=1, save_weights_only=True, period=1
)

When arriving at the line where the weights are loaded, I receive the following error:

ValueError: Trying to create optimizer slot variable under the scope for 
tf.distribute.Strategy 
(<tensorflow.python.distribute.distribute_lib._DefaultDistributionStrategy
 object at 0x7fb6047c3c50>), which is different from the scope used for the original 
variable (MirroredVariable:{

Any help would be appreciated

Answer 1

According to this model-saving APIs create variables (which should be distributed variables). Try creating/calling the callbacks from within the strategy scope.