A training step updates a model after it sees one batch. It makes predictions, compares them with the batch’s labels to get a loss, computes how the loss changes with the model’s weights, and updates those weights to reduce the loss. section 2.5.2
One batch of handwritten digits
For an MNIST classifier, pass a batch of images through the model, calculate sparse categorical cross-entropy against the digit labels, and average the per-image losses. TensorFlow’s GradientTape obtains gradients of that average with respect to the model’s weights; an update function then applies them. The step returns the average loss so training progress can be observed. section 2.5.2 section 2.5.3
one_training_step <- function(model, images_batch, labels_batch) {
with(tf$GradientTape() %as% tape, {
predictions <- model$call(images_batch)
per_sample_losses <-
op_sparse_categorical_crossentropy(labels_batch, predictions)
average_loss <- op_mean(per_sample_losses)
})
gradients <- tape$gradient(average_loss, model$weights)
update_weights(gradients, model$weights)
average_loss
}
The complete R function places prediction and loss calculation inside GradientTape, then obtains gradients and updates the weights. section 2.5.2
The same sequence has a few important choices:
- Update size is controlled by the learning rate. Subtracting gradient times learning rate is a simple update; a Keras optimizer usually handles the update instead. section 2.5.2
- A step can use a small batch or all training samples. The linear-classifier example uses all 2,000 points per step: each step costs more computation but incorporates every sample. section 3.3.2
- In PyTorch, loss$backward() populates parameter gradients. Apply updates without tracking their gradients, then clear the stored gradients; otherwise they accumulate across steps. section 3.4.2
Repeating a step over every batch makes an epoch; repeating epochs makes a training loop. To customize a Keras fit() run without replacing that loop, override a model’s train_step() method and return named metric values. A custom step can use losses and metrics supplied to compile(); JAX additionally requires explicit handling of model and optimizer state. section 2.5.3 section 7.4.4 section 7.4.5 section 7.4.2
