Follow up of this issue https://github.com/keras-team/keras/issues/20979

The given solutions works on the same running notebook.

def get_config(self):
    return {
        "name": self.name,
        "input_shape": self.input_shape[1:],
        "filters": self.filters,
        "activation": self.activation,
    }

But later if model.keras is loaded it gives

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-2-36406e627be8>](https://localhost:8080/#) in <cell line: 0>()
     41 # model.save('new.keras') # ok
     42 
---> 43 loaded_model = keras.saving.load_model("new (1).keras")

5 frames
[/usr/local/lib/python3.11/dist-packages/keras/src/saving/serialization_lib.py](https://localhost:8080/#) in _retrieve_class_or_fn(name, registered_name, module, obj_type, full_config, custom_objects)
    801             return obj
    802 
--> 803     raise TypeError(
    804         f"Could not locate {obj_type} '{name}'. "
    805         "Make sure custom classes are decorated with "

TypeError: Could not locate class 'DummyModel'. Make sure custom classes are decorated with `@keras.saving.register_keras_serializable()`. Full object config: {'module': None, 'class_name': 'DummyModel', 'config': {'input_shape': [28, 28, 1], 'filters': [16, 32], 'activation': 'relu'}, 'registered_name': 'Custom>DummyModel', 'build_config': {'input_shape': None}}

Full code.

import numpy as np

import keras
from keras import layers
keras.__version__ # 3.8

@keras.saving.register_keras_serializable()
class DummyModel(keras.Model):
    def __init__(
        self,
        *,
        input_shape=(28, 28, 1),
        filters=[16, 32],
        activation='relu',
        **kwargs,
    ):
        input_spec = keras.layers.Input(shape=input_shape)
        x = input_spec
        x = layers.Conv2D(filters[0], 3, activation=activation)(x)
        x = layers.Conv2D(filters[1], 3, activation=activation)(x)
        x = layers.MaxPooling2D(3)(x)
        x = layers.Conv2D(filters[1], 3, activation=activation)(x)
        x = layers.Conv2D(filters[0], 3, activation=activation)(x)
        x = layers.GlobalMaxPooling2D()(x)
        super().__init__(inputs=input_spec, outputs=x, **kwargs)

        self.filters = filters
        self.activation = activation

    def get_config(self):
        config = {
                "input_shape": self.input_shape[1:],
                "filters": self.filters,
                "activation": self.activation,
            }
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)

a = np.ones((1, 28, 28, 1), dtype=np.float32); print(a.shape)
model = DummyModel()
output = model(a)
print(output.shape)

model.save('new.keras') # ok
loaded_model = keras.saving.load_model("new.keras") # ok

How to reproduce

  1. Generate the .keras file (in colab)
  2. Download the keras file.
  3. Disconnect and create an new session (in colab) and upload the file
  4. Run: loaded_model = keras.saving.load_model("new.keras") # error

Comment From: innat

cc. @mattdangerw @sonali-kumari1

Comment From: mattdangerw

@innat I believe this is expected. Keras does not pickle or make any attempt to save your python custom objects -- pickle is not safe, and we want to keep our saved model simply config and weights. Instead the user is responsible for bringing the definitions for any custom objects through one of the following (this will need to be done in the session that is loading the object):

  • Registering custom objects (preferred),
  • Passing custom objects directly when loading, or
  • Using a custom object scope

See https://keras.io/guides/serialization_and_saving/

If you want a saved method that captures all weights, config, and the full computation graph, take a look at model.export() to the tf or onnx format. Both these have a safe format that captures all computation (with the caveat of course that the original python classes will be gone). Or you could roll your own solution with pickle, but given safety and reliability issues with pickle I don't think Keras will attempt to make a save format that contains arbitrary user python code.