Repository: keras-team/keras

Issue:
Error deserializing saved model after Keras bump 3.6 -> 3.7
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 795df4ef63566ec869fe2512373f4346e1e02746.

The function was serialized without the package name prior to 795df4ef63566ec869fe2512373f4346e1e02746, 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')


Discussion:
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.

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 0c2bdff313f7533f0d7e6670a906102cc2fb046d and the associated pull request https://github.com/keras-team/keras/pull/20406 in the changelog for Keras 3.7: https://github.com/keras-team/keras/compare/v3.6.0...v3.7.0.
> 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](https://colab.research.google.com/gist/Surya2k1/cc7e7d5fc45d7c6c0a2cb3755aaeb68a/20806_3-6v.ipynb).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?

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](https://colab.research.google.com/gist/mehtamansi29/396ea6a0e6293aa1f383bea098883886/20806-error-deserializing-saved-model-after-keras-bump-3-6-3-7.ipynb) 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](https://keras.io/guides/customizing_saving_and_serialization/#build-and-compile-saving-customization) you can find more details about @keras.saving.register_keras_serializable(package="my_custom_package"). 
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: https://github.com/keras-team/keras/blob/734cd03b9413e65dea392e25c420676dd04f080a/keras/src/losses/__init__.py#L139 
and activation functions: https://github.com/keras-team/keras/blob/734cd03b9413e65dea392e25c420676dd04f080a/keras/src/activations/__init__.py#L95

In Keras 3.6, custom loss functions were saved using only the function name as key, ignoring the package in the decorator. The PR https://github.com/keras-team/keras/pull/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.
> Here in the [gist](https://colab.research.google.com/gist/mehtamansi29/396ea6a0e6293aa1f383bea098883886/20806-error-deserializing-saved-model-after-keras-bump-3-6-3-7.ipynb) 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?
> 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](https://keras.io/guides/serialization_and_saving/#custom-objects). 


> 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](https://github.com/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](https://colab.research.google.com/gist/mehtamansi29/0e37cd67a076e90cc31b4a7f05aff7bd/copy-of-20806_3-6v2.ipynb) shows model is loading properly with custom object code. 
@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](https://keras.io/guides/serialization_and_saving/#custom-objects), 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")
```
> 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.
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 https://github.com/keras-team/keras/pull/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()))
```
Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSdHag0RVFS7UXzZkKcsFCKOcX8raCupKK9RHSlYxp5U8lSJbQ/viewform?entry.492125872=Yes&entry.243948740=https://github.com/keras-team/keras/issues/20806">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSdHag0RVFS7UXzZkKcsFCKOcX8raCupKK9RHSlYxp5U8lSJbQ/viewform?entry.492125872=No&entry.243948740=https://github.com/keras-team/keras/issues/20806">No</a>




Patch:
diff --git a/keras/src/saving/serialization_lib.py b/keras/src/saving/serialization_lib.py
index 48c70808b405..e01d22d52728 100644
--- a/keras/src/saving/serialization_lib.py
+++ b/keras/src/saving/serialization_lib.py
@@ -781,6 +781,16 @@ def _retrieve_class_or_fn(
                 if obj is not None:
                     return obj
 
+            # Workaround for serialization bug in Keras <= 3.6 whereby custom
+            # functions would only be saved by name instead of registered name,
+            # i.e. "name" instead of "package>name". This allows recent versions
+            # of Keras to reload models saved with 3.6 and lower.
+            if ">" not in name:
+                separated_name = ">" + name
+                for custom_name, custom_object in custom_objects.items():
+                    if custom_name.endswith(separated_name):
+                        return custom_object
+
         # Otherwise, attempt to retrieve the class object given the `module`
         # and `class_name`. Import the module, find the class.
         package = module.split(".", maxsplit=1)[0]


Test patch:
diff --git a/keras/src/saving/serialization_lib_test.py b/keras/src/saving/serialization_lib_test.py
index 4891d961089d..06c22acfc327 100644
--- a/keras/src/saving/serialization_lib_test.py
+++ b/keras/src/saving/serialization_lib_test.py
@@ -8,6 +8,7 @@
 import keras
 from keras.src import ops
 from keras.src import testing
+from keras.src.saving import object_registration
 from keras.src.saving import serialization_lib
 
 
@@ -343,6 +344,29 @@ def test_layer_sharing(self):
         serialized, deserialized, reserialized = self.roundtrip(func)
         self.assertLen(deserialized.layers, 3)
 
+    def test_keras36_custom_function_reloading(self):
+        @object_registration.register_keras_serializable(package="serial_test")
+        def custom_registered_fn(x):
+            return x**2
+
+        config36 = {
+            "module": "builtins",
+            "class_name": "function",
+            "config": "custom_registered_fn",
+            "registered_name": "function",
+        }
+        obj = serialization_lib.deserialize_keras_object(config36)
+        self.assertIs(obj, custom_registered_fn)
+
+        config = {
+            "module": "builtins",
+            "class_name": "function",
+            "config": "serial_test>custom_registered_fn",
+            "registered_name": "function",
+        }
+        obj = serialization_lib.deserialize_keras_object(config)
+        self.assertIs(obj, custom_registered_fn)
+
 
 @keras.saving.register_keras_serializable()
 class MyDense(keras.layers.Layer):