qubvel / segmentation_models

Segmentation models with pretrained backbones. Keras and TensorFlow Keras.
MIT License
4.77k stars 1.03k forks source link

Error when loading a previously saved model #466

Closed twVolc closed 3 years ago

twVolc commented 3 years ago

Using tensorflow version 2.2.0

I'm getting the following error when using model = tf.keras.models.load_model(model_path) on a model previously saved with model.save(). I've tried saving the model in both Tensorflow Native format and hdf5 and I get the same error:

Traceback (most recent call last):
  File "C:/Users/tw9616/Documents/PostDoc/ML/Python project/roc_curves.py", line 13, in <module>
    model = tf.keras.models.load_model(model_path)
  File "C:\Users\tw9616\Anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\keras\saving\save.py", line 184, in load_model
    return hdf5_format.load_model_from_hdf5(filepath, custom_objects, compile)
  File "C:\Users\tw9616\Anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\keras\saving\hdf5_format.py", line 193, in load_model_from_hdf5
    model.compile(**saving_utils.compile_args_from_training_config(
  File "C:\Users\tw9616\Anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\keras\saving\saving_utils.py", line 215, in compile_args_from_training_config
    loss = _deserialize_nested_config(losses.deserialize, loss_config)
  File "C:\Users\tw9616\Anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\keras\saving\saving_utils.py", line 255, in _deserialize_nested_config
    return deserialize_fn(config)
  File "C:\Users\tw9616\Anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\keras\losses.py", line 1831, in deserialize
    return deserialize_keras_object(
  File "C:\Users\tw9616\Anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\keras\utils\generic_utils.py", line 392, in deserialize_keras_object
    raise ValueError('Unknown ' + printable_module_name + ':' + object_name)
ValueError: Unknown loss function:binary_crossentropy_plus_jaccard_loss

It seems to be an issue with the segmentation-models built-in loss function binary_crossentropy_plus_jaccard_loss which I used for training my model. Have I been saving my models in the wrong way for segmentation-models or is this a bug that needs fixing? Presumably I could just save the model weights and rebuild the model that way, but it seems easier for deployment to be able to load the model with the single line of code tf.keras.models.load_model().

fidhaph commented 3 years ago

you need to mention the custom objects while loading the model. in your case mention "binary_crossentropy_plus_jaccard_loss"

twVolc commented 3 years ago

@fidhaph brilliant thank you, that worked a treat.

Just to clarify in case anyone else is having this issue, you pass the loss function in a dictionary to the custom_objects argument of load_model. The dictionary key is the string name of the function: tf.keras.models.load_model(model_path, custom_objects={'binary_crossentropy_plus_jaccard_loss': sm.losses.bce_jaccard_loss})