Overriding train_step

Setup#

pip install tensorflow numpy

Outputs verified with TensorFlow 2.21 (Keras 3.12), CPU, seed 3. This is TensorFlow-backend code: inside the step we use tf.GradientTape — meet it properly in GradientTape.

A train step with built-in input noise#

The model below augments its inputs with small random noise during training only — a regulariser baked into the step, invisible at prediction time.

import numpy as np
import tensorflow as tf
import keras
from keras import layers

keras.utils.set_random_seed(3)


class NoisyTrainer(keras.Model):
    """A model that adds small input noise during training only —
    a data augmentation baked into the training step."""

    def train_step(self, data):
        x, y = data
        x = x + tf.random.normal(tf.shape(x), stddev=0.05)

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

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

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


inputs = keras.Input(shape=(2,))
outputs = layers.Dense(1, activation="sigmoid")(layers.Dense(16, activation="relu")(inputs))
model = NoisyTrainer(inputs, outputs)

x = np.array([[0., 0.], [0., 1.], [1., 0.], [1., 1.]], dtype="float32")
y = np.array([[0.], [1.], [1.], [0.]], dtype="float32")   # XOR

model.compile(optimizer=keras.optimizers.Adam(0.05),
              loss="binary_crossentropy", metrics=["accuracy"])
history = model.fit(x, y, epochs=400, batch_size=4, verbose=0)

print("loss at epoch 1:  ", round(history.history["loss"][0], 4))
print("loss at epoch 400:", round(history.history["loss"][-1], 4))
print("predictions:", np.round(model.predict(x, verbose=0).ravel(), 2))

loss at epoch 1:   0.7381
loss at epoch 400: 0.0006
predictions: [0. 1. 1. 0.]

The walkthrough#

The five standard moves are all visible. Forward pass under the tape, compute_loss, tape.gradient, apply_gradients, metric updates. This is exactly what stock Keras runs when you do not override — the override is your chance to insert changes between the moves. Ours is one line: the noise addition before the tape.

self.compute_loss(y=y, y_pred=y_pred) asks the model to apply whatever loss you passed to compile. Using it, rather than hard-coding a loss, keeps the override reusable — compile still decides the objective.

self(x, training=True) matters. Layers like Dropout and BatchNorm behave differently in training; the flag is how they know. Forget it and those layers run in inference mode during training — the same trap as train and eval mode in PyTorch.

The metrics loop keeps the scoreboard honest: the loss metric tracks the loss value, every other metric compares truth against prediction. The returned dict is what callbacks and the progress bar display — same contract as stock Keras.

Noise made XOR slower, not impossible. Clean XOR trains in far fewer epochs. With noise, inputs like 0.03, 0.97 stand in for 0, 1, so the model must learn a tolerant boundary. That tolerance is the regularising point.

Common mistakes#

Overriding when compile-and-fit would do. Sample weighting, class weighting, gradient clipping (clipnorm=1.0 on any optimizer), and learning-rate schedules all have stock support. The override is for schemes that genuinely change the step's structure.

Computing the forward pass outside the tape. Move y_pred = self(x) above with tf.GradientTape(): and gradients come back as None — the tape never saw the computation. Everything the gradient must flow through goes inside the with block.

Forgetting test_step. evaluate and validation_split use a separate method, test_step. Override only train_step and validation still runs the stock path — correct here (we want no noise at validation), but wrong if your change affects the loss itself. Override both when the loss logic changes.

Returning nothing. The returned dict feeds logging and callbacks. Return nothing and EarlyStopping has no numbers to watch.

Try it yourself#

Add gradient clipping inside the step: grads = [tf.clip_by_norm(g, 1.0) for g in grads] before apply_gradients. Confirm training still converges. Then check the stock alternative — delete the line and pass clipnorm=1.0 to the Adam constructor instead. Same effect, no override; background in gradient clipping.

What to learn next#

• GradientTape — the recording device the custom step leans on.
• GAN — the architecture that made custom steps famous.
• Gradient clipping — the most common thing people add to a step.