LayerGenius | Saving and Loading Keras 3 Models: Formats, Custom Layers, and a Round-Trip Test

A model that trains is half a deliverable. A model that loads, byte-for-byte equivalent, in a different process six months later is the other half. Saving and loading is where many Keras 3 migrations quietly break, so this post covers the formats, the failure modes, and a round-trip test that catches them.

Why this deserves a test

Keras 3 saves models in the .keras format: a zip archive holding config.json, metadata.json, and model.weights.h5. The architecture is rebuilt from config, then weights are loaded by layer path. Two things follow:

Any custom class that can't serialize its config can't be rebuilt.
Weights match by structure, so renaming or reordering layers can break a load without touching the code that trains.

Legacy SavedModel directories and .h5 files are different paths. model.save("path.keras") is the default to standardize on. For TF Serving or other SavedModel consumers, use model.export("dir"), which writes an inference artifact, not a re-trainable model.

Three artifacts, three purposes

Artifact	Call	Use it for
`.keras` archive	`model.save` / `keras.saving.load_model`	Resuming, fine-tuning, handing a model to another team
Weights only	`model.save_weights("m.weights.h5")`	Checkpoints when code defines the architecture
Export artifact	`model.export`	Serving; no Python class needed

Pick on purpose. Mixing them up is the most common source of "works on my machine" load errors we see in reviews.

Make custom layers serializable

import keras

@keras.saving.register_keras_serializable(package="layergenius")
class ScaledResidual(keras.layers.Layer):
    def __init__(self, units, scale=0.1, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        self.scale = scale
        self.dense = keras.layers.Dense(units)

    def call(self, x):
        return x + self.scale * self.dense(x)

    def get_config(self):
        cfg = super().get_config()
        cfg.update({"units": self.units, "scale": self.scale})
        return cfg

Common mistakes:

Forgetting an argument in get_config. The model loads, with the default value. This is the silent one.
Creating sublayers in call instead of __init__ or build. Their weights are not tracked.
Passing a lambda or a raw function as an argument. Register the function, or refer to it by name.

Lambda layers and `safe_mode`

keras.saving.load_model defaults to safe_mode=True, which refuses to deserialize arbitrary Python code from Lambda layers. Don't switch it off to make an error go away. Turn it off only for files you produced yourself, and better, replace the Lambda with a registered layer. Treat model files from outside your organization like pickles: untrusted until proven otherwise.

Weights-only checkpoints in training

For long runs, use keras.callbacks.ModelCheckpoint with save_weights_only=True and a .weights.h5 filename. Checkpoints are smaller and faster. Save the full .keras file once at the end of training, along with the git SHA and the config that built it.

If optimizer state matters for resuming, save the full model. Weights-only restarts reset optimizer moments, so the first epochs after resume look worse than the last ones before it. That's expected, not a bug.

Cross-backend loading

A .keras file saved under one backend generally loads under another, because weights are stored as backend-neutral arrays. Two caveats:

Custom layers must use keras.ops, not backend-specific calls, or the load succeeds and the forward pass fails.
Numerics differ slightly across backends. Compare outputs with a tolerance, not equality.

Measure before you promise portability. The backend benchmarking post covers how to compare backends on your own model.

The round-trip test

Put this in CI. It takes seconds on a small config.

import numpy as np, keras, tempfile, os

def test_roundtrip(build_model, example_batch):
    model = build_model()
    expected = model.predict(example_batch, verbose=0)

    with tempfile.TemporaryDirectory() as d:
        path = os.path.join(d, "m.keras")
        model.save(path)
        restored = keras.saving.load_model(path)

    got = restored.predict(example_batch, verbose=0)
    np.testing.assert_allclose(expected, got, rtol=1e-5, atol=1e-6)
    assert len(model.weights) == len(restored.weights)

Add two more cases once this passes: load in a fresh Python process (no custom classes imported by accident), and load under each backend you claim to support.

A checklist before hand-off

Every custom class is registered and has a complete get_config.
No Lambda layers in a model that leaves the team.
The round-trip test passes in a fresh process.
The artifact is stored with Keras and backend versions, the training config, and a hash of the weights file.
Serving gets an export artifact, verified separately. See the export post.

Trade-offs

The .keras format is stable for a given Keras major version but not a long-term archive. If a model must be loadable in three years, keep the code and a pinned environment next to it, and re-run the round-trip test when you upgrade. That costs an hour per upgrade. Finding out a model won't load during an incident costs far more.

Saving and Loading Keras 3 Models: Formats, Custom Layers, and a Round-Trip Test