ValueError: Unable to restore custom object of type _tf_keras_metric currently.

[Abdelrahman Hassanein]

I've trained a Unet model based on the multiclass segmentation (camvid) example and I've an issue trying to load it...

I've used the following code and callbacks:

# create model
model = sm.Unet(BACKBONE, classes=n_classes, activation=activation)

# define optomizer
optim = keras.optimizers.Adam(LR)

# loss
focal_loss = sm.losses.BinaryFocalLoss() if n_classes == 1 else sm.losses.CategoricalFocalLoss()

# metrics
metrics = [sm.metrics.IOUScore(threshold=0.5), sm.metrics.FScore(threshold=0.5)]

# compile
model.compile(optim, focal_loss, metrics)

I've changed the callback to save model in savedmodel format instead of .h5, my callbacks looked like this:

callback_save = tf.keras.callbacks.ModelCheckpoint(
    wandb.config.checkpoint_name, monitor='val_loss', 
    verbose=1, save_best_only=True,
    save_weights_only=False, mode='min', 
    save_freq='epoch')

callbacks = [
    callback_save,
    keras.callbacks.ReduceLROnPlateau(),
    WandbCallback()
]

when try to load saved model and passing the custom_objects as follows:

keras.models.load_model(r'checkpoints\baseline_best_001', custom_objects={'categorical_focal_loss': sm.losses.CategoricalFocalLoss,
                                                        'iou_score': sm.metrics.IOUScore,
                                                        'f1_score': sm.metrics.FScore})

I get the error

ValueError: Unable to restore custom object of type _tf_keras_metric currently. Please make sure that the layer implements `get_config`and `from_config` when saving. In addition, please use the `custom_objects` arg when calling `load_model()`.