Subclassing keras.Model

Two methods, two jobs

A subclassed model is a Python class that inherits from keras.Model and splits its work across exactly two methods. __init__() is the parts list: you instantiate every layer the model owns and store it on self. call() is the assembly: it receives an input tensor and threads it through those layers to produce an output. Keras takes over from there — once you've defined call(), the instance behaves like any built-in model, so compile(), fit(), evaluate(), and save() all work unchanged (see the Code block).

Why the split matters

The separation isn't stylistic. Layers created in __init__() are instantiated once when you write model = MyModel() — and their weights are tracked from that moment, which is how Keras knows what to optimize. call() runs on every forward pass, once per batch. Put a layer construction inside call() by mistake and you'll silently create fresh, untrained weights on each step; put tensor math inside __init__() and there's no tensor to operate on yet. The training argument that call() receives is what lets Dropout and BatchNormalization switch between their train-time and inference-time behavior — thread it through to every sub-layer that cares.

Code

A minimal subclassed keras.Model·python

import keras
from keras import layers

class MyModel(keras.Model):
    def __init__(self):
        super().__init__()
        self.dense1 = layers.Dense(128, activation="relu")
        self.dropout = layers.Dropout(0.3)
        self.dense2 = layers.Dense(10, activation="softmax")

    def call(self, inputs, training=False):
        x = self.dense1(inputs)
        x = self.dropout(x, training=training)
        return self.dense2(x)

model = MyModel()
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
model.fit(x_train, y_train, epochs=5)

Subclassing keras.Model

Two methods, two jobs

Why the split matters

Code

External links

Exercise

Progress

Comments 0