Custom Training Step

Keep fit(), change the step

There's a tier of customization between "model.fit() as-is" and "write the whole training loop by hand." Override train_step() and you replace just the per-batch logic — the gradient computation — while keeping everything fit() gives you for free: the progress bar, callbacks like EarlyStopping and ModelCheckpoint, validation passes, and distribution strategy. This is the right tool for GANs, knowledge distillation, curriculum learning, or any setup where a vanilla loss-and-backprop step isn't enough (see the Code block).

The contract train_step() must honor

train_step(self, data) receives one batch, and Keras expects two things back. First, you must actually update the weights inside the method — compute the loss, derive gradients against self.trainable_variables, and apply them through self.optimizer. Second, you must return a dict of metric names to values; that dict is exactly what fit() prints in the progress bar and feeds to every callback. Use self.compute_loss() rather than calling the loss function directly — it folds in any regularization and lets compile(loss=...) keep working. Get the dict wrong and the callbacks see nothing; skip the optimizer apply and the model trains on nothing.

Code

Overriding train_step() while keeping fit()·python

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

        # Forward pass with gradient tracking
        y_pred = self(x, training=True)
        loss = self.compute_loss(y=y, y_pred=y_pred)

        # Compute and apply gradients
        gradients = self.optimizer.compute_gradients(loss, self.trainable_variables)
        self.optimizer.apply(gradients)

        # Update 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}

Keep fit(), change the step

The contract train_step() must honor

Code

External links

Exercise

Progress

Comments 0