Customizing what fit() does

GANs, contrastive learning and many research setups need custom steps, and this keeps all of fit()'s conveniences.

fit() calls model.train_step(data) for every batch. Overriding it gives you control over the step while keeping callbacks, progress bars, distribution and evaluate/predict. Because gradients are computed differently on each backend, the step is written in that backend's style:

TensorFlow: compute the loss inside a tf.GradientTape, get gradients with tape.gradient, and apply with self.optimizer.apply(grads, self.trainable_variables).
PyTorch: call self.zero_grad(), compute the loss, loss.backward(), collect v.value.grad for each trainable variable, and apply under torch.no_grad().
JAX: steps are stateless: train_step(self, state, data) receives and returns variable values, and uses self.stateless_call and self.optimizer.stateless_apply.

In every version, use self.compute_loss(x=x, y=y, y_pred=y_pred) to get the loss configured in compile (including regularisation losses), and update self.metrics. Return a dict of metric results; those become the logs callbacks see.

Override test_step the same way for custom evaluation. If you need several optimizers (as in a GAN), override compile to accept them. For full control without fit(), write a training loop from scratch with the same building blocks.

Example

import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import numpy as np
import keras
import tensorflow as tf

class CustomModel(keras.Model):
    def train_step(self, data):
        x, y = data
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss = self.compute_loss(x=x, y=y, y_pred=y_pred)
        grads = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply(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((32,))
model = CustomModel(inputs, keras.layers.Dense(1)(inputs))
model.compile(optimizer="adam", loss="mse", metrics=["mae"])
model.fit(np.random.rand(256, 32), np.random.rand(256, 1), epochs=2, verbose=0)

Remember

Override train_step to customise the step but keep fit()
Gradient code is backend-specific; JAX steps are stateless
Use self.compute_loss and update self.metrics; return their results

Watch out for

Forgetting to update metrics in the custom step, so fit() logs stale or zero values.

Check yourself

What is the main advantage of overriding train_step instead of writing a full training loop?

Official docs: keras.io/guides/custom_train_step_in_tensorflow/

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks