skip to content

Why does Keras load_model fail to locate your custom layer, and how do you fix it?

level: middleimportance: must knowfreq 68%

answer

  1. the archive stores names, not code
  2. the decorator must have actually run
  3. import the module before loading
  4. a registry or an explicit mapping argument

basics

~20 s

The .keras archive stores a class name and its constructor arguments, never the code. On load, Keras must map that name to a real Python class; if it cannot, the load fails. Fix it with @keras.saving.register_keras_serializable or by passing custom_objects.

solid answer

~40 s

Saving serializes your layer as a name plus a config dict. `keras.saving.load_model` then has to find a class for that name, and a fresh interpreter that never imported your module cannot. Two fixes: decorate the class with `@keras.saving.register_keras_serializable(package="MyPkg")`, which puts it in a global registry - but the module must still be imported before the load, so the decorator actually runs - or pass the mapping explicitly, `keras.saving.load_model(path, custom_objects={"MyLayer": MyLayer})`. `keras.saving.custom_object_scope` applies the same mapping to a block of code. Registration is the durable choice for library code; custom_objects is fine in a script. The rule covers custom losses, metrics, activations and initializers too.

code

python · 19 lines
python
import keras

@keras.saving.register_keras_serializable(package="MyPkg")
class Scale(keras.layers.Layer):
    def __init__(self, factor=2.0, **kwargs):
        super().__init__(**kwargs)
        self.factor = factor

    def call(self, x):
        return x * self.factor

    def get_config(self):
        config = super().get_config()
        config.update({"factor": self.factor})
        return config

m = keras.Sequential([keras.layers.Input((4,)), Scale(3.0)])
m.save("scaled.keras")
m2 = keras.saving.load_model("scaled.keras")

go deeper

for a junior

Know that a saved model stores class names rather than code, and that reloading a model with your own layer needs either the register_keras_serializable decorator or a custom_objects mapping.

for a middle

Explain the resolution path: config.json holds the name, load_model looks it up in the registry or the custom_objects dict, and the decorator only helps if its module was imported.

for a senior

Bring the production angle: import model packages in the serving entrypoint, add a save-and-reload smoke test for every custom layer, and namespace registrations with package to avoid collisions.

for a principal

Treat serializability as a review standard. Custom code that cannot round-trip is an artifact you cannot redeploy, so decide where registration is mandatory and where teams should use built-in layers instead.

## Why the failure exists at all The `.keras` v3 archive is deliberately code-free. `config.json` records, for each object, a class name (with a module path and any registered package prefix) and the keyword arguments needed to construct it. That is what makes a saved model inspectable and easy to move around - and it is exactly why loading a model containing your own classes needs help. When `keras.saving.load_model` walks the config and reaches an entry whose class it cannot resolve, it raises while trying to locate that class, telling you to register the class as serializable or pass it through `custom_objects`. ## Fix 1: register the class ``` @keras.saving.register_keras_serializable(package="MyPkg") class MyLayer(keras.layers.Layer): ... ``` The decorator adds the class to Keras's global custom-object registry under `MyPkg>MyLayer`, and saving records that same string. Loading then resolves it with no extra arguments - the ergonomic win, because callers of your package need not know the class exists. The catch that trips people: a decorator only runs when its module is imported. If the serving process loads the archive without ever importing the module that defines `MyLayer`, the registry is empty and the load fails just as before. Import your model package before calling `load_model`. The `package` argument is worth setting. It namespaces the name, so two projects that each define `Attention` do not collide in one registry. ## Fix 2: pass custom_objects ``` model = keras.saving.load_model( "model.keras", custom_objects={"MyLayer": MyLayer, "my_loss": my_loss} ) ``` The dictionary keys are the names Keras recorded in the config - the bare class name by default, or the registered `package>name` string if the class was registered when it was saved. Values are the actual classes or functions. `keras.saving.custom_object_scope({...})` is the same mapping applied to everything loaded inside a `with` block, handy when several loads happen together. The two mechanisms compose: custom_objects wins for names it covers, and the registry fills in the rest. ## It is not only layers Anything whose implementation is yours goes through the same door: a subclassed Loss or Metric, a plain-function loss passed to compile, a custom activation, constraint, initializer or regularizer. A model that trains perfectly can fail to reload purely because of a one-line loss function - which is why teams register their custom objects as a habit rather than after the first outage. ## Two related traps First, resolving the class is only half the job. Keras calls it with the recorded config, so a class whose `get_config` omits a constructor argument is re-created with the wrong argument, or raises a TypeError for a missing one. Registration makes the class findable; a complete config makes it re-creatable. Second, `load_model` defaults to `safe_mode=True`, which refuses to deserialize `Lambda` layers because their bodies are stored as marshalled Python bytecode that would execute on load. If you own the file and truly need it, `safe_mode=False` allows it; the better answer is to replace the Lambda with a registered Layer subclass. The habit that avoids all of this: after writing a custom layer, immediately save and reload a tiny model containing it as a unit test. Serialization bugs surface at reload time, which in production is the worst possible moment to discover them.

  • The class is decorated with register_keras_serializable but the serving process still cannot load it. Why?
    Because the decorator only executes when its module is imported. A serving entrypoint that imports keras and calls load_model, without importing your model package, leaves the registry empty. Import the module explicitly at startup, or pass custom_objects. This is the most common cause of a load that works in the training notebook and fails in production.
  • What are the keys of the custom_objects dictionary?
    The names as they appear in the saved config: the bare class or function name by default, or the package-qualified string such as MyPkg>MyLayer if the object was registered with a package when it was saved. If you are unsure, unzip the archive and read config.json - the string you need is there verbatim.
  • Does this apply to a custom loss function too?
    Yes. Losses, metrics, activations, initializers, constraints and regularizers are all serialized by name and resolved the same way. A one-line custom loss is enough to break a reload, so register it or pass it in custom_objects. Loading with compile=False sidesteps a loss-only problem when you only need inference.

saying these in an interview costs you the question

  • Believing save() embeds the layer's source code
  • Registering the class but never importing its module
  • Passing the layer instance name instead of the class
  • Assuming custom losses are exempt from registration
  • Reaching for safe_mode=False as the first fix

context

resolving links…