Writing custom layers

Custom layers are how you package new ideas so they work with fit, saving and every backend.

A custom layer subclasses `keras.layers.Layer` and usually implements:

__init__(self, ..., **kwargs): store hyperparameters; call super().__init__(**kwargs).
build(self, input_shape): create weights with self.add_weight(shape, initializer, trainable=True, name=...). Creating weights here, rather than in __init__, lets the layer adapt to the input size it first sees.
call(self, inputs, training=None): the computation, written with keras.ops for portability.
get_config(self): return the constructor arguments merged with super().get_config(), so the layer can be re-created with from_config when a saved model is loaded.

Decorate the class with @keras.saving.register_keras_serializable(package="...") so load_model can find it without custom_objects.

Other useful hooks: self.add_loss(value) for activity regularisation, non-trainable weights (trainable=False) for state like moving averages (update them with assign), and compute_output_shape if Keras cannot infer it.

Nested layers assigned as attributes are tracked automatically, so their weights appear in layer.weights. Test a new layer by calling it on a NumPy batch, checking output shapes, and round-tripping it through get_config/from_config.

Example

import numpy as np
import keras
from keras import ops

@keras.saving.register_keras_serializable(package="demo")
class ScaledDense(keras.layers.Layer):
    def __init__(self, units, scale=1.0, **kwargs):
        super().__init__(**kwargs)
        self.units, self.scale = units, scale
    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                                 initializer="glorot_uniform", name="w")
        self.b = self.add_weight(shape=(self.units,), initializer="zeros", name="b")
    def call(self, x):
        return self.scale * (ops.matmul(x, self.w) + self.b)
    def get_config(self):
        return {**super().get_config(), "units": self.units, "scale": self.scale}

layer = ScaledDense(4, scale=0.5)
print(tuple(layer(np.ones((2, 3), dtype="float32")).shape))   # (2, 4)
print([tuple(w.shape) for w in layer.weights])                # [(3, 4), (4,)]
clone = ScaledDense.from_config(layer.get_config())
print(clone.units, clone.scale)                               # 4 0.5

Remember

Create weights in build with add_weight
Write call with keras.ops for portability
Implement get_config and register the class for saving

Watch out for

Creating weights in call, which makes new weights on every call.

Check yourself

Why create weights in build() rather than __init__?

Official docs: keras.io/guides/making_new_layers_and_models_via_subclassing/

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks