Subclassing keras.Model

Subclassing gives full Python flexibility for research architectures that are awkward as static graphs.

For architectures with custom control flow, subclass `keras.Model`:

Create layers in __init__ and assign them to attributes; Keras tracks their weights automatically, including layers stored in lists or dicts.
Implement call(self, inputs, training=False) with the forward computation, passing training on to layers that need it (Dropout, BatchNorm).
Call super().__init__(**kwargs) first so arguments like name work.

A subclassed model is used like any other: compile, fit, evaluate, predict. Weights are created on the first call, so call the model on a sample batch (or call build) before summary().

Trade-offs compared with the functional API:

More flexible: loops, conditionals and arbitrary Python inside call.
Less introspection: Keras cannot plot the layer graph or check shapes before running.
Saving to .keras works, but to reload you must implement get_config (if __init__ takes arguments) and register the class.

A good pattern is to build reusable blocks by subclassing keras.layers.Layer and assemble them with the functional API, getting flexibility where you need it and introspection everywhere else. With the PyTorch backend, a Keras model is also a torch.nn.Module.

Example

import numpy as np
import keras
from keras import layers

class ResidualMLP(keras.Model):
    def __init__(self, units=64, n_classes=10, **kwargs):
        super().__init__(**kwargs)
        self.inp = layers.Dense(units, activation="relu")
        self.block = layers.Dense(units, activation="relu")
        self.drop = layers.Dropout(0.2)
        self.head = layers.Dense(n_classes)

    def call(self, x, training=False):
        h = self.inp(x)
        h = h + self.drop(self.block(h), training=training)   # residual connection
        return self.head(h)

model = ResidualMLP()
print(tuple(model(np.zeros((2, 20), dtype="float32")).shape))   # (2, 10)  builds weights
print(len(model.trainable_weights))                              # 6
model.compile(optimizer="adam",
              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True))

Remember

Layers in __init__, computation in call(inputs, training=False)
Weights exist after the first call
Flexible but less introspectable than functional models

Watch out for

Ignoring the training argument, so dropout stays active or inactive in the wrong phase.

Check yourself

Why might model.summary() fail right after constructing a subclassed model?

Official docs: keras.io/guides/making_new_layers_and_models_via_subclassing/

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks