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?

[mehtamansi29]

Hi @Surya2k1 and @cloudbopper -

Here while saving the model with serialization with custom object need to use get_config and from_config method.
Here in the gist I tried to create sample model with @keras.saving.register_keras_serializable in keras3.6.0 and loading the model in same keras3.7.0 working fine. But on restart the session custom objects while deserialize the model is not saved after restart session. So need to re-register those custom objects again in keras3.6.0 and also loading model fine with upgrade keras version.
Here you can find more details about @keras.saving.register_keras_serializable(package="my_custom_package").

[cloudbopper]

Ok, it turns out that the deserialization only fails for custom loss functions, not activation functions, when loading models saved in 3.6 in 3.8. I modified @Surya2k1's gist to show this behavior: https://colab.research.google.com/gist/cloudbopper/0efab83644d74000ebac5a5025f2c37e/20806_3-6v2.ipynb. The last cell output shows the error.

The reason appears to be serialization differences between loss functions:

keras/keras/src/losses/__init__.py

Line 139 in 734cd03

return serialization_lib.serialize_keras_object(loss)

and activation functions:

keras/keras/src/activations/__init__.py

Line 95 in 734cd03

fn_config["config"] = object_registration.get_registered_name(

In Keras 3.6, custom loss functions were saved using only the function name as key, ignoring the package in the decorator. The PR #20406 that was merged into 3.7 fixed this, but also removed fallback code that was allowing deserialization to work in 3.6, leading to the error in 3.7 and 3.8.

The serialization code for activation functions adds the package to the key using get_registered_name so it avoids the issue.

[Surya2k1]

Here in the gist I tried to create sample model with @keras.saving.register_keras_serializable in keras3.6.0 and loading the model in same keras3.7.0 working fine. But on restart the session custom objects while deserialize the model is not saved after restart session. So need to re-register those custom objects again in keras3.6.0 and also loading model fine with upgrade keras version.

Hi @mehtamansi29 , That's what we have been discussing. The question is, for custom objects to reload, Is it the only way to add custom object code in order to reload the model? or Is there any other way? Suppose I have a saved model with custom objects and don't have access to custom object code can I able to reload the model? Could you please confirm on it?

[mehtamansi29]

That's what we have been discussing. The question is, for custom objects to reload, Is it the only way to add custom object code in order to reload the model? or Is there any other way? Suppose I have a saved model with custom objects and don't have access to custom object code can I able to reload the model? Could you please confirm on it?

@Surya2k1 - While working with custom layers, activations and losses in model saving and reloading, custom objects is preferable for model saving(serializing) and reloading(deserializing) with subclass layer.
For that while saving model must need to define get_config() and from_config() method. Saved .keras file is lightweight and does not store the Python code for custom objects.
For reload the model requires access to the definition of any custom objects via 3 methods: Registering custom objects,
Passing custom objects directly when loading and Using a custom object scope.
More details about custom object and those 3 methods you can find here.

Ok, it turns out that the deserialization only fails for custom loss functions, not activation functions, when loading models saved in 3.6 in 3.8. I modified @Surya2k1's gist to show this behavior: https://colab.research.google.com/gist/cloudbopper/0efab83644d74000ebac5a5025f2c37e/20806_3-6v2.ipynb. The last cell output shows the error.

@cloudbopper - Here in your gist, while reloading the model in keras3.8.0, you are not using custom object code properly at here. reloaded_model = keras.models.load_model("custom_model.keras").
Attached gist shows model is loading properly with custom object code.

[cloudbopper]

@mehtamansi29 it looks like you replaced

reloaded_model = keras.models.load_model(
    "custom_model.keras"
)

with

custom_objects = {
    "MyLayers": CustomLayer,
    "custom_fn": custom_fn,
    "my_loss_fn": my_loss_fn
}

reloaded_model = keras.models.load_model(
    "custom_model.keras",custom_objects=custom_objects
)

for loading the model in both Keras versions.

However, I'm just using the recommended method (registering custom objects using the @keras.saving.register_keras_serializable decorator) from the documentation you shared, the first method you mentioned to @Surya2k1:

This is the preferred method, as custom object registration greatly simplifies saving and loading code. Adding the @keras.saving.register_keras_serializable decorator to the class definition of a custom object registers the object globally in a master list, allowing Keras to recognize the object when loading the model.

You can see in their example that registering custom objects avoids having to pass in custom_objects while loading the model:

# Now, we can simply load without worrying about our custom objects.
reconstructed_model = keras.models.load_model("custom_model.keras")

[Surya2k1]

You can see in their example that registering custom objects avoids having to pass in custom_objects while loading the model:

Agreed. No need to pass custom_objects arguments with the decorator.

When we pass custom_objects explicitly to load_model, For example in this case, it will be duplicated and custom_objects will become double like below,with Package name and without function names:

{'MyLayers': <class '__main__.CustomLayer'>, 'MyLayers>CustomLayer': <class '__main__.CustomLayer'>, 'custom_fn': <function custom_fn at 0x0000029D776D440>, 'custom_fn>custom_fn': <function custom_fn at 0x0000029D776D440>, 'custom_fn>my_loss_fn': <function my_loss_fn at 0x0000029DA6E2C00>, 'my_loss_fn': <function my_loss_fn at 0x0000029DA6E2C00>}

You can see without explicit passing of custom objects it will have only package_name>fun_name as key and when we pass it explicitly one more item will be added as we are passing only fun_name in custom_object it will also get appended to custom_objects and hence it works since the saved model in keras 3.6 has only fun_name.

This can be a workaround for loading saved model in 3.6V with custom_loss in Versions >3.7. If it is OK then documentation needs to be updated accordingly.

I also have a probable fix in hand, but eventually it will become redundant(In Future versions when 3.6 become too old to use), which can avoid loading of custom_objects explicitly.

[cloudbopper]

Thanks! Yes, I was able to work around it by passing custom_objects while loading the model - just wanted to make sure it's addressed for other users. The easiest fix might be to just re-add the fallback code in keras/src/saving/saving_lib_test.py that was removed in #20406, although it's a bit of a hack:

            # Retrieval of registered custom function in a package
            filtered_dict = {
                k: v
                for k, v in custom_objects.items()
                if k.endswith(full_config["config"])
            }
            if filtered_dict:
                return next(iter(filtered_dict.values()))

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No