Subclassing keras.layers.Layer

The layer is the real workhorse

Subclassing a Layer is far more common than subclassing a whole Model — it's the recommended building block, because a custom layer drops straight into Sequential, Functional, or another subclassed model without ceremony. A layer has three methods worth knowing: __init__() stores configuration (how many units, which activation), build() creates the weights, and call() runs the forward pass (see the Code block).

Why weights live in build(), not init()

The interesting one is build(input_shape). You could create weights in __init__(), but then you'd have to know the input's feature dimension up front and hardcode it — brittle the moment you reuse the layer somewhere else. build() is Keras's answer: it runs automatically the first time the layer sees real data, receives the concrete input_shape, and sizes the weight matrix from input_shape[-1]. This is lazy weight creation, and it's why you can write MyDense(64) and have it adapt to whatever feeds it. Each weight is registered with add_weight(), which is what makes Keras track it as a trainable variable.

get_config() closes the loop on serialization: it returns the constructor arguments as a plain dict so Keras can rebuild the layer from a saved file. Skip it and your custom layer won't round-trip through model.save() / load_model().

Code

A custom Dense layer with build() and get_config()·python

class MyDense(keras.layers.Layer):
    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        # Lazy weight creation — called on first use
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="glorot_uniform",
            trainable=True,
            name="kernel",
        )
        self.b = self.add_weight(
            shape=(self.units,),
            initializer="zeros",
            trainable=True,
            name="bias",
        )

    def call(self, inputs):
        return keras.ops.matmul(inputs, self.w) + self.b

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

Subclassing keras.layers.Layer

The layer is the real workhorse

Why weights live in build(), not init()

Code

External links

Exercise

Progress

Comments 0

The layer is the real workhorse

Why weights live in build(), not __init__()

Code

External links

Exercise

Progress

Why weights live in build(), not init()