Fix Keras model loading failures with custom layers in Transformer ASR by using register_keras_serializable and proper get_config

Learn why Keras models with custom layers fail to load and how to fix it: register_keras_serializable, implement get_config, and reload with custom_objects.

Saving and reloading Keras models that use custom layers can be unexpectedly tricky. If you have followed the Transformer ASR tutorial and ended up with a model that trains and saves fine, yet fails to load with a ValueError about missing variables, you are likely facing an issue with serialization of custom objects. Here is a clear walkthrough of the failure mode and how to fix it properly.

Reproducing the problem

After training a speech recognition model for one epoch, the model is saved without errors:

net.save("asr_epoch1.keras")

When attempting to load the model back with its custom components, the following pattern is used:

catalog = {
    'TokenEmbedding': LexemeEmbedding,
    'SpeechFeatureEmbedding': AudioFeatureBlock,
    'TransformerEncoder': StackEncoder,
    'TransformerDecoder': StackDecoder,
    'Transformer': Seq2SeqTransformer,
    'CustomSchedule': WarmupSchedule
}
restored = keras.models.load_model("asr_epoch1.keras", custom_objects=catalog, compile=False)

But loading fails with an error like this:

ValueError: A total of 51 objects could not be loaded. Example error message for object <Dense name=dense_65, built=True>:

Layer 'dense_65' expected 2 variables, but received 0 variables during loading. Expected: ['kernel', 'bias']

List of objects that could not be loaded:
[<Dense name=dense_65, built=True>, <Embedding name=embedding_10, built=True>, <Embedding name=embedding_11, built=True>, <EinsumDense name=key, built=True>, <EinsumDense name=attention_output, built=True>, <EinsumDense name=query, built=True>, <EinsumDense name=value, built=True>, <Dense name=dense_63, built=True>, <Dense name=dense_64, built=True>, <LayerNormalization name=layer_normalization_63, built=True>, <LayerNormalization name=layer_normalization_64, built=True>, <LayerNormalization name=layer_normalization_65 .......

The model itself is assembled and trained like this:

net = Seq2SeqTransformer(
    num_hid=200,
    num_head=2,
    num_feed_forward=400,
    target_maxlen=max_target_len,
    num_layers_enc=4,
    num_layers_dec=1,
    num_classes=34,
)
opt = keras.optimizers.Adam(lr)
net.compile(optimizer=opt, loss=loss_obj)
hist = net.fit(train_ds, validation_data=eval_ds, callbacks=[progress_cb], epochs=1)

Saving to H5 was also attempted:

net.save("asr_epoch1.h5")

but that failed as well when loading.

What is going on

The error is raised during deserialization, when Keras tries to reconstruct layers and load their weights. The message “expected N variables, but received 0” points to custom components not being properly set up for serialization. Ensuring that custom layers are registered for Keras serialization and that they return their initialization parameters in get_config is required so Keras can rebuild them correctly while loading. Crucially, these definitions must be in place before the model is built and saved; otherwise, the saved artifact won’t carry the necessary metadata to restore the variables.

The fix

Register each custom component with @keras.saving.register_keras_serializable and make get_config return the exact arguments used in __init__. Then rerun the notebook from a clean state, rebuild the model, retrain, save, and only then load it with custom_objects.

Here is an example of a serializable audio embedding block used inside the ASR model:

@keras.saving.register_keras_serializable(package="components")
class AudioFeatureBlock(layers.Layer):
    def __init__(self, hid_units=64, max_len=100):
        super().__init__()
        self.hid_units = hid_units
        self.max_len = max_len
        self.f1 = keras.layers.Conv1D(
            hid_units, 11, strides=2, padding="same", activation="relu"
        )
        self.f2 = keras.layers.Conv1D(
            hid_units, 11, strides=2, padding="same", activation="relu"
        )
        self.f3 = keras.layers.Conv1D(
            hid_units, 11, strides=2, padding="same", activation="relu"
        )

    def get_config(self):
        return {
            "hid_units": self.hid_units,
            "max_len": self.max_len,
        }

    def call(self, inp):
        z = self.f1(inp)
        z = self.f2(z)
        return self.f3(z)

After registering custom layers and returning init parameters in get_config, rerun the notebook and retrain the model so the saved file reflects the updated serialization setup. You can then load the model like this:

catalog = {
    'TokenEmbedding': LexemeEmbedding,
    'SpeechFeatureEmbedding': AudioFeatureBlock,
    'TransformerEncoder': StackEncoder,
    'TransformerDecoder': StackDecoder,
    'Transformer': Seq2SeqTransformer,
    'CustomSchedule': WarmupSchedule
}
net2 = keras.models.load_model("asr_epoch1.keras", custom_objects=catalog, compile=False)

Why this matters

When building ASR systems or any deep model with custom components, you often iterate in notebooks. Execution order and state persistence matter: if the serialization hooks are added after the model was created, the saved artifact isn’t consistent and won’t reload. Registering components and providing a complete get_config before any training ensures that the model file is portable and reloadable across sessions and environments.

Practical wrap-up

Define and register every custom layer or utility you use in the model graph. Make get_config return all constructor arguments. Restart the environment, rebuild the model, train, and save to the .keras format. Load with custom_objects and compile=False when needed. If you get stuck, reduce the code to a minimal snippet that builds, saves, and loads the model—this helps isolate the exact point of failure quickly.