Module 4 — Custom layers and Model subclassing

The two previous APIs assemble existing layers. This module goes one level down: writing the layer itself, then taking back control of the training loop. You need this less often than you might think, but when the need arises there is no way around it.

One layer, three methods

Every custom layer inherits from keras.layers.Layer and spreads across three methods with clearly separated roles.

from tensorflow import keras
import tensorflow as tf

class ScaledDense(keras.layers.Layer):
    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units            # hyperparameters only

    def build(self, input_shape):
        self.kernel = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="glorot_uniform",
            trainable=True,
            name="kernel",
        )
        self.scale = self.add_weight(
            shape=(), initializer="ones", trainable=True, name="scale",
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.kernel) * self.scale

__init__ receives the hyperparameters and nothing else. build creates the weights, and it runs once, on the first batch. call describes the computation.

The reason build exists is simple: at __init__ time the input dimension is unknown. build receives it via input_shape, which is what lets you write Dense(64) without ever stating how many inputs arrive. Creating weights in __init__ works if you hard-code the dimension, but you lose that flexibility and the layer becomes unusable elsewhere.

add_weight, not tf.Variable

A tf.Variable declared directly inside a layer is not reliably tracked: it may not appear in layer.trainable_weights, may not be saved, and may never receive a gradient. add_weight registers it with Keras. The symptom of getting this wrong is a weight that never moves, with no message at all.

Training mode must be forwarded

Some layers behave differently depending on whether you are training or predicting. They must receive that information and pass it on.

class RegularisedBlock(keras.layers.Layer):
    def __init__(self, units, rate=0.3, **kwargs):
        super().__init__(**kwargs)
        self.dense = keras.layers.Dense(units, activation="relu")
        self.dropout = keras.layers.Dropout(rate)

    def call(self, inputs, training=None):
        h = self.dense(inputs)
        return self.dropout(h, training=training)

Omitting training=training in the Dropout call produces a nasty bug: dropout stays active during evaluation and prediction. Validation scores become noisy and systematically pessimistic, predictions change from one call to the next, and nothing points to the cause. The same trap applies to BatchNormalization, whose statistics must not update outside training.

Serialise so you can reload

Saving a model that contains a custom layer is not enough: Keras must know how to rebuild that layer on load.

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

get_config returns the __init__ arguments as a dictionary. Without it, keras.models.load_model fails with an unknown-object error. With it, plus the registration decorator, reloading is transparent:

@keras.saving.register_keras_serializable(package="inskillml")
class RegularisedBlock(keras.layers.Layer):
    ...

It is a piece of housekeeping, but it is the one that decides whether your model is deployable — the subject of module 10.

Subclass Model when the flow is dynamic

Subclassing keras.Model moves the architecture definition into imperative Python code.

class AdaptiveClassifier(keras.Model):
    def __init__(self, num_classes, **kwargs):
        super().__init__(**kwargs)
        self.trunk = keras.layers.Dense(128, activation="relu")
        self.head = keras.layers.Dense(num_classes, activation="softmax")

    def call(self, inputs, training=None):
        h = self.trunk(inputs)
        if training:
            h = tf.nn.dropout(h, rate=0.2)
        return self.head(h)

The gain is total freedom: conditionals, loops whose iteration count depends on the data, recursive calls. The cost is real and often underestimated. No graph exists before the first batch, so model.summary() stays silent until then, shape errors surface only at runtime, and the model no longer serialises as simply.

The decision rule is clear: stay functional as long as the architecture is a fixed graph. Subclassing earns its place when the flow depends on the data at runtime, which is rare outside research.

Override train_step rather than the whole loop

This is the most useful point in the module. When the standard loop does not fit, the common reaction is to write a full loop with GradientTape — and in doing so lose callbacks, progress bars, distributed training and history. There is a far more economical entry point.

class CustomLossModel(keras.Model):
    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            prediction = self(x, training=True)
            loss = self.compute_loss(x=x, y=y, y_pred=prediction)

        gradients = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        for metric in self.metrics:
            if metric.name != "loss":
                metric.update_state(y, prediction)
        return {m.name: m.result() for m in self.metrics}

train_step receives one batch and returns a dictionary of metrics. All the rest of the Keras infrastructure keeps working: fit, the callbacks from module 6, TensorBoard from module 7, the distribution from module 9. This is the right level of intervention for gradient clipping, adversarial training, gradient accumulation across batches, or a loss that depends on more than the input-output pair.

training=True in the self(x, ...) call is not optional: it is what activates dropout and the normalisation statistics update.

Key takeaways

• A custom layer separates hyperparameters in __init__, weights in build and computation in call; build receives the input shape, which makes the layer reusable.
• Create weights with add_weight, not tf.Variable, or they escape Keras tracking and never train.
• Forward training=training to sub-layers; omitting it leaves dropout active at prediction time, with pessimistic scores and no error message.
• To change training behaviour, override train_step rather than writing a full loop: you keep fit, callbacks, TensorBoard and distribution.

Next module: tf.data pipelines, because a fast model fed too slowly is still slow.