Keras Custom Models

The subclassing API is Keras's most flexible option — writing a custom tf.keras.Model subclass, directly parallel to PyTorch's custom nn.Module pattern from nn.Module.

The Pattern — call() Instead of forward()

import tensorflow as tf
from tensorflow.keras import layers

class MyModel(tf.keras.Model):
    def __init__(self):
        super().__init__()   # ALWAYS call this first -- same importance as PyTorch's 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)
        if training:
            x = self.dropout(x)   # explicit training-mode check -- an important Keras-specific pattern
        return self.dense2(x)

model = MyModel()
output = model(x, training=True)   # calling the instance directly, exactly like PyTorch's model(x)

The Key Difference From PyTorch: The training Argument

Instead of PyTorch's separate model.train()/model.eval() mode-switching (from nn.Module), Keras's call() method receives an explicit training boolean argument, letting the model's forward logic branch directly based on it — a subtly different design for achieving the same goal (different behavior during training versus inference for layers like dropout and batch normalization).

A ResNet-Style Block, Subclassed

class ResidualBlock(tf.keras.Model):
    def __init__(self, filters):
        super().__init__()
        self.conv1 = layers.Conv2D(filters, 3, padding='same', activation='relu')
        self.conv2 = layers.Conv2D(filters, 3, padding='same')

    def call(self, inputs):
        x = self.conv1(inputs)
        x = self.conv2(x)
        return layers.ReLU()(x + inputs)   # the residual connection, expressed directly in Python

Common Mistakes

• Forgetting to pass and check the training argument for layers like dropout that need to behave differently between training and inference — without explicitly handling it, dropout may always (or never) apply, regardless of the actual mode.
• Forgetting super().__init__() — analogous to the same PyTorch mistake, this breaks Keras's internal layer-tracking machinery.

Interview Relevance

Q: "How does Keras's subclassing API handle training-mode-dependent behavior (like dropout), compared to PyTorch's approach?" Keras passes an explicit training boolean argument into the call() method, which the model's code can check directly to branch its behavior. PyTorch instead tracks a persistent training/eval mode flag on the module itself (set via model.train()/model.eval()), which layers like nn.Dropout check internally without needing an explicit argument passed through every call.

Practice Question

What is the Keras subclassing API's equivalent of PyTorch's forward() method?

Related DL Notes

• nn.Module
• Keras Functional API
• Keras Training