Overriding train_step()

~10 min · custom-train

Level 0Keras Apprentice

0 XP0/97 lessons0/20 achievements

0/120 XP to next level120 XP to go0% complete

The sweet spot of customization

Overriding train_step() is the first rung you should reach for: you replace exactly one method — the per-batch forward/backward logic — and inherit everything else. The progress bar, callbacks, validation pass, and distribution strategy all keep working because fit() still drives the loop; it just calls your step instead of the built-in one.

The four moves inside the step

Every train_step() does the same four things: unpack the batch, run a forward pass to get y_pred, compute the loss with self.compute_loss(), then compute and apply gradients. You finish by updating the metric objects and returning a name → value dict — that dict is exactly what fit() prints in the progress bar. The code below shows the TensorFlow-backend version; the gradient mechanics are where the backend leaks through (see the warning).

Code

Override train_step() (TensorFlow backend)·python

import tensorflow as tf
import keras

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

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

        # Compute and apply gradients (TF-native)
        grads = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(
            zip(grads, self.trainable_variables)
        )

        # Update and return metrics
        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}

# Still use fit()!
model = CustomModel(...)
model.compile(optimizer="adam", loss="mse")
model.fit(x_train, y_train, epochs=10)

External links

Exercise

Subclass keras.Model with a train_step that adds an L2 regularization term to the gradient manually. Train one epoch on MNIST and verify the regularizer appears in the loss.