A training step updates a model using one batch of inputs and targets. It makes predictions, measures their loss against the targets, computes how the loss changes with the trainable weights, and adjusts those weights to reduce the loss. This is the unit of work repeated throughout training. section 2.5.2 section 2.5.3
One batch of handwritten digits
For a batch of images and labels, run the model on the images, compute sparse categorical crossentropy against the labels, and average the losses. Differentiate that average with respect to the model weights, then apply the gradients. The returned loss tells you how this batch went; it is not itself the weight update. section 2.5.2
def one_training_step(model, images_batch, labels_batch): with tf.GradientTape() as tape: predictions = model(images_batch) loss = ops.sparse_categorical_crossentropy(labels_batch, predictions) average_loss = ops.mean(loss) gradients = tape.gradient(average_loss, model.weights) update_weights(gradients, model.weights) return average_loss
A TensorFlow implementation records prediction and loss operations with GradientTape, differentiates the average loss, and updates the weights. section 2.5.2
A simple update subtracts gradient × learning rate from each weight: the gradient gives the direction of increasing loss, and the learning rate sets the size of the move in the opposite direction. In practice, use an optimizer such as Keras SGD rather than writing that update by hand. Automatic differentiation supplies gradients without manually deriving backpropagation through operations such as softmax and crossentropy. section 2.5.2
The surrounding loop has a different job. It calls the step once per batch to make an epoch, then repeats epochs. A custom step controls what happens to one batch; a training loop controls which batches are processed and for how long. section 2.5.3
The four operations stay recognizable across backends, but gradient and state handling change:
- TensorFlow records the forward pass inside tf.GradientTape(), then asks the tape for gradients and applies them with an optimizer. section 7.4.2
- PyTorch calls loss.backward(), reads weight gradients, applies the update, and clears gradients with model.zero_grad(). section 7.4.2
- JAX uses jax.value_and_grad() and a stateless step that returns updated weights. A Keras step using JAX also carries non-trainable and optimizer state. section 3.5.4 section 7.4.2
learning_rate = 0.1 @jax.jit #1 def training_step(inputs, targets, W, b): loss, grads = grad_fn((W, b), inputs, targets) #2 grad_wrt_W, grad_wrt_b = grads W = W - grad_wrt_W * learning_rate #3 b = b - grad_wrt_b * learning_rate #3 return loss, W, b #4
The small JAX classifier makes the stateless pattern concrete: compute loss and gradients, update W and b, and return all three values. section 3.5.4
When only the per-batch algorithm needs changing, override keras.Model.train_step() and continue using fit(). Return a dictionary of metric names and values; this retains fit() features such as callbacks, performance optimizations, and distributed training support. If the loss and metrics were supplied to compile(), self.compute_loss() incorporates the configured loss and applicable regularization losses, while self.metrics provides the metrics to update and report. section 7.4.4 section 7.4.5
