Addons: metrics causes error while restoring model from checkpoints which are saved using keras callbacks.

[ramsane]

System information

• Linux ( Google colab)
• TensorFlow version 2.2.0
• TensorFlow-Addons version and how it was installed (source or binary):
• Python version:
• Is GPU used? yes

Describe the bug

I have used callbacks for model checkpointing, And for the metrics, I have used f1-score from tfaddons. It is giving error while loading the checkpoints

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()`.

Code to reproduce the issue

 # compile the model
 macro_f1_score = tfa.metrics.F1Score(num_classes=10,average='macro', name='macro_f1_score')
 model2.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=1e-3),
               loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
               metrics=[tf.keras.metrics.CategoricalAccuracy(), macro_f1_score])
  # callback
  tf.keras.callbacks.ModelCheckpoint(filepath='./ckpt/model2/{epoch}/',
                                     monitor='val_macro_f1_score', verbose=1,  mode='min',
                                     save_best_only=True, save_freq='epoch')
# ... .. .. 
# restoring
checkpoints = [checkpoint_dir + "/" + name for name in os.listdir(checkpoint_dir)]
if checkpoints:
  latest_checkpoint = max(checkpoints, key=os.path.getctime)
  print("Restoring from", latest_checkpoint)
  model2 = tf.keras.models.load_model(latest_checkpoint)
else:
  print('Training from scratch')

This is not the whole file, I just pasted the snippets that are necessary. Main things are using callbacks for model checkpoints, use tfa.metrics in the metrics while compiling the model, and restoring the model. We can easily reproduce this if these things are like this.

Other info / logs
I've tried including that as a custom object with key and value as macro_f1_score`, but the same error.

[bhack]

This is not the whole file, I just pasted the snippets that are necessary.

Could you share a very dummy minimal but complete example to reproduce this? It is easier for us to quickly copy, run and debug with less overheads.

Thanks

[ramsane]

Sorry for the delay. I've taken a sample notebook from the TensorFlow documentation and modified few things.

https://colab.research.google.com/drive/1A86gxNw-GibMj5Ochi-imZG1ZCRdVkL-?usp=sharing

With this, you should be able to reproduce the error.

[bhack]

Have you tried with:

model = tf.keras.models.load_model(latest_checkpoint, custom_objects={'F1Score': macro_f1_score})

[ramsane]

It worked. Thanks.

I have tried with custom_objects={'macro_f1_score'': macro_f1_score} and custom_objects={'f1_score': macro_f1_score}.

When I searched for it, They mentioned that we have to provide the function name. But here are initializing F1Score and and calling on top of it.

Is there any documentation on what name to use in custom_objects OR would you let us know so that it will be helpful for me and others as well.

[bhack]

https://www.tensorflow.org/api_docs/python/tf/keras/models/load_model

It already tell you to use class or function string name. In this case Is a class

[bhack]

Please open a Documentation issue or a Documentation PR directly if you think that It could be improved.