tensorflow: ValueError: Unknown metric function: CustomMetric using custom metrics when loading tf saved model type with tf.keras.models.load_model

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

System information

  • Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
  • OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 18.04
  • TensorFlow installed from (source or binary): binary
  • TensorFlow version (use command below): 2.0.0
  • Python version: 3.7

Describe the current behavior ValueError: Unknown metric function: CustomMetric occurs when trying to load a tf saved model using tf.keras.models.load_model with a custom metric. If you look at the code for load_model, it is clear the load_model currently ignores the custom_objects dict for the tf saved model format.

Describe the expected behavior load_model loads the custom metric successfully either just implicitly or through the custom_objects dict.

Code to reproduce the issue

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.metrics import Metric
import numpy as np

class CustomMetric(Metric):
  def __init__(self,
               name='score',
               dtype=tf.float32):
    super(CustomMetric, self).__init__(name=name)
    self.true_positives = self.add_weight(
        'true_positives',
        shape=[10],
        initializer='zeros',
        dtype=self.dtype)


  def update_state(self, y_true, y_pred, sample_weight=None):
    pass

  def result(self):
    return 0

  def get_config(self):
    """Returns the serializable config of the metric."""
    config = {}
    base_config = super(CustomMetric, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

  def reset_states(self):
    self.true_positives.assign(np.zeros(self.num_classes), np.float32)
    self.weights_intermediate.assign(
        np.zeros(self.num_classes), np.float32)
            
inputs = keras.Input(shape=(784,), name='digits')
x = layers.Dense(64, activation='relu', name='dense_1')(inputs)
x = layers.Dense(64, activation='relu', name='dense_2')(x)
outputs = layers.Dense(10, activation='softmax', name='predictions')(x)
model = keras.Model(inputs=inputs, outputs=outputs, name='3_layer_mlp')

model.compile(loss='sparse_categorical_crossentropy', optimizer=tf.keras.optimizers.Adam(lr=.001), metrics=[CustomMetric()])

model.save("model/", save_format='tf')

new_model = keras.models.load_model('model/',  tf.keras.models.load_model ={'score': CustomMetric})

Other info / logs

Traceback (most recent call last):
  File "/home/sentim/Website/model_prediction/test_load_saved_model.py", line 46, in <module>
    new_model = keras.models.load_model('model/', custom_objects={'score': CustomMetric})
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/save.py", line 150, in load_model
    return saved_model_load.load(filepath, compile)
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py", line 93, in load
    model._training_config))  # pylint: disable=protected-access
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 356, in compile
    self._cache_output_metric_attributes(metrics, weighted_metrics)
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 1901, in _cache_output_metric_attributes
    metrics, self.output_names, output_shapes, self.loss_functions)
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_utils.py", line 813, in collect_per_output_metric_info
    metric_name = get_metric_name(metric, is_weighted)
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_utils.py", line 987, in get_metric_name
    metric = metrics_module.get(metric)
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/keras/metrics.py", line 2857, in get
    return deserialize(identifier)
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/keras/metrics.py", line 2851, in deserialize
    printable_module_name='metric function')
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/keras/utils/generic_utils.py", line 180, in deserialize_keras_object
    config, module_objects, custom_objects, printable_module_name)
  File "/home/sentim/anaconda3/envs/py37/lib/python3.7/site-packages/tensorflow_core/python/keras/utils/generic_utils.py", line 165, in class_and_config_for_serialized_keras_object
    raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
ValueError: Unknown metric function: CustomMetric

About this issue

  • Original URL
  • State: closed
  • Created 5 years ago
  • Comments: 22 (8 by maintainers)

Most upvoted comments

The models saved in h5 format seem to work fine, the issue is about models saved with SavedModel format (as explained here https://www.tensorflow.org/guide/saved_model)

What is working is setting the compile flag to False and then compiling it on its own e.g.:

import tensorflow as tf

# Save Keras Model as SavedModel (Keras model has some custom objects e.g. custom loss function)
tf.saved_model.save(my_keras_model, EXPORT_DIR)
...
# Load the model and compile on its own (working)
loaded_model = tf.keras.model.load(EXPORT_DIR, custom_objects={"custom_loss": my_custom_loss}, compile=False)
model.compile(optimizer=my_optimizer, loss=my_custom_loss)
# Load the model while also loading optimizer and compiling (failing with "Unkown loss function: my_custom_loss")
loaded_model = tf.keras.model.load(EXPORT_DIR, custom_objects={"custom_loss": my_custom_loss}) # compile is set to True by default

The point is:

  1. The default way of loading models fails if there are custom objects involved.
  2. By compiling yourself you are setting up a new optimizer instead of loading the previously trained models optimizer weights.

Moreover I already submited a PR that would fix this: https://github.com/tensorflow/tensorflow/pull/34048. But it seems nobody bothers about it : /

Here is a workaround for the meantime:

from tensorflow.python.saved_model import loader_impl
from tensorflow.python.keras.saving.saved_model import load as saved_model_load

loader_impl.parse_saved_model(load_path)
model = saved_model_load.load(load_path, custom_objects={"custom_metric": custom_metric})

Here is a workaround for the meantime:

from tensorflow.python.saved_model import loader_impl
from tensorflow.python.keras.saving.saved_model import load as saved_model_load

loader_impl.parse_saved_model(load_path)
model = saved_model_load.load(load_path, custom_objects={"custom_metric": custom_metric})

not working at keras 2.3.1, tf 2.0.0

TypeError: load() got an unexpected keyword argument 'custom_objects'

@jvishnuvardhan While it does work in the h5 format, if I have saved a model to the tf format, I cannot load the model to resave it to the h5 format later (since I can’t load the model in the first place), so ultimately this is still an issue that needs to be addressed.

I have this problem loading an .h5 model on TF 2.3.0.

same issue here, when you save the model in tf format, you can’t re-load the model with custom_objects, this should be fixed.