keras.models.load_model() gives ValueError

Question

I have saved the trained model and the weights as below.

model, history, score = fit_model(model, train_batches, val_batches, callbacks=[callback])
model.save('./model')
model.save_weights('./weights')

Then I tried to get the saved model as the following way

if __name__ == '__main__':
  model = keras.models.load_model('./model', compile= False,custom_objects={"F1Score": tfa.metrics.F1Score})
  test_batches, nb_samples = test_gen(dataset_test_path, 32, img_width, img_height)
  predict, loss, acc = predict_model(model,test_batches, nb_samples)
  print(predict)
  print(acc)
  print(loss)

But it gives me an error. What should I do to overcome this?

Traceback (most recent call last):
  File "test_pro.py", line 34, in <module>
    model = keras.models.load_model('./model',compile= False,custom_objects={"F1Score": tfa.metrics.F1Score})
  File "/home/dcs2016csc007/.local/lib/python3.8/site-packages/tensorflow/python/keras/saving/save.py", line 212, in load_model
    return saved_model_load.load(filepath, compile, options)
  File "/home/dcs2016csc007/.local/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 138, in load
    keras_loader.load_layers()
  File "/home/dcs2016csc007/.local/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 379, in load_layers
    self.loaded_nodes[node_metadata.node_id] = self._load_layer(
  File "/home/dcs2016csc007/.local/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 407, in _load_layer
    obj, setter = revive_custom_object(identifier, metadata)
  File "/home/dcs2016csc007/.local/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 921, in revive_custom_object
    raise ValueError('Unable to restore custom object of type {} currently. '
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()`.

Answer 1

Looking at the source code for Keras, the error is raised when trying to load a model with a custom object:

def revive_custom_object(identifier, metadata):
  """Revives object from SavedModel."""
  if ops.executing_eagerly_outside_functions():
    model_class = training_lib.Model
  else:
    model_class = training_lib_v1.Model

  revived_classes = {
      constants.INPUT_LAYER_IDENTIFIER: (
          RevivedInputLayer, input_layer.InputLayer),
      constants.LAYER_IDENTIFIER: (RevivedLayer, base_layer.Layer),
      constants.MODEL_IDENTIFIER: (RevivedNetwork, model_class),
      constants.NETWORK_IDENTIFIER: (RevivedNetwork, functional_lib.Functional),
      constants.SEQUENTIAL_IDENTIFIER: (RevivedNetwork, models_lib.Sequential),
  }
  parent_classes = revived_classes.get(identifier, None)

  if parent_classes is not None:
    parent_classes = revived_classes[identifier]
    revived_cls = type(
        compat.as_str(metadata['class_name']), parent_classes, {})
    return revived_cls._init_from_metadata(metadata)  # pylint: disable=protected-access
  else:
    raise ValueError('Unable to restore custom object of type {} 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()`.'
                     .format(identifier))

The method will only work fine with the custom objects of the types defined in revived_classes. As you can see, it currently only works with input layer, layer, model, network, and sequential custom objects.

In your code, you pass an tfa.metrics.F1Score class in the custom_objects argument, which is of type METRIC_IDENTIFIER, therefore, not supported (probably because it doesn't implement the get_config and from_config functions as the error output says):

keras.models.load_model('./model', compile=False, custom_objects={"F1Score": tfa.metrics.F1Score})

It's been a while since I last worked with Keras but maybe you can try and follow what was proposed in this other related answer and wrap the call to tfa.metrics.F1Score in a method. Something like this (adjust it to your needs):

def f1(y_true, y_pred):
    metric = tfa.metrics.F1Score(num_classes=3, threshold=0.5)
    metric.update_state(y_true, y_pred)
    return metric.result()

keras.models.load_model('./model', compile=False, custom_objects={'f1': f1})