Error deserializing saved model after Keras bump 3.6 -> 3.7

[cloudbopper]

I'm encountering this error while loading a model saved with Keras 3.6 in Keras 3.7 and 3.8:

TypeError: Could not locate function 'attention_loss'. Make sure custom classes are decorated with `@keras.saving.register_keras_serializable()`. Full object config: {'module': 'builtins', 'class_name': 'function', 'config': 'attention_loss', 'registered_name': 'function'}

The decorator for the custom function:

@tf.keras.utils.register_keras_serializable("attention_loss")
def attention_loss(y_actual, y_predicted):

I narrowed down the changes causing the error to 795df4e.

The function was serialized without the package name prior to 795df4e, so the removal of the fallback code means that the serialized function key can no longer be used to retrieve the registered object (key 'attention_loss>attention_loss')

[Surya2k1]

Hi @cloudbopper ,

Please try the below steps:

1.If you have custom loss function code then you can copy the code before reloading the model.
2.Then pass the argument dict custom_objects to load_model() like below.

@keras.saving.register_keras_serializable("attention_loss")
def attention_loss(y_actual, y_predicted):
   """Implementation here"""

reloaded_model = keras.models.load_model(
    "custom_model.keras",
    custom_objects={"attention_loss": attention_loss},
)

This should work probably. I don't think its related to the version. For example if you have a model with custom classes or functions then save the model and restart the session and try to reload the model in same version without adding code to custom object like mentioned above. The problem will be replicated.

The root cause seems Keras not saving complete custom object implementation except class/function name and signature. For the function body/implementation it depends on the storage ID of this object locally. When a session restarted it will loose this storage. Hence we need to add the custom objects before reloading the model from .keras .

May be providing option for serializing custom_objects during model.save() will be a better option so that users have the flexibility to save entire model along with custom objects.

@mehtamansi29 , Correct me If I am mistaken in any thing I said.

[cloudbopper]

Yes, it works by adding custom_objects={"attention_loss": attention_loss}, to load_model but it shouldn't be required since the function was registered using @tf.keras.utils.register_keras_serializable decorator when the model was saved. The docs state that using the decorator is the preferred method https://keras.io/guides/serialization_and_saving/. I've gotten around it using custom_objects but I wanted to flag the issue since function de/serialization behavior shouldn't be backward incompatible between minor versions, or should at least be documented so others don't have to spend time debugging it.

I can confirm that saving the model with the function registered using the decorator @tf.keras.utils.register_keras_serializable in Keras 3.6, and then loading it in a new session is working in Keras 3.6 and no longer working in Keras 3.7.

You can see the responsible merge commit 0c2bdff and the associated pull request #20406 in the changelog for Keras 3.7: v3.6.0...v3.7.0.

[Surya2k1]

I can confirm that saving the model with the function registered using the decorator @tf.keras.utils.register_keras_serializable in Keras 3.6, and then loading it in a new session is working in Keras 3.6 and no longer working in Keras 3.7.

Hi @cloudbopper , Actually I have tested a custom model with Keras 3.6v, saved it and restarted the session and tried to reload the session. Got error in reloading the model if not added custom object code in new session. Please refer to attached colab gist.Same is true for other version also like 3.5v and 3.8v.

So what I observed is without adding custom code I can't reload a saved model with custom objects(even with serialization registry), can't be reloaded in a new session or with a new version without explicitly adding the code for custom objects before reloading.

Could you please verify the colab gist that I have attached and confirm whether you have not observed similar behavior?