Training: connect prediction, loss and an update

Training changes a model using a learning signal. For a neural network, a typical step computes predictions, measures their loss, differentiates that loss and applies an optimizer update.

Inference uses the current parameters to produce an output. Training adds the machinery needed to decide how those parameters should change.

The complete update in one function

This TensorFlow step accepts a model returning class logits, an optimizer, inputs and integer class labels. It averages the example losses into one objective.

import tensorflow as tf

def train_step(model, optimizer, inputs, labels):
    with tf.GradientTape() as tape:
        logits = model(inputs, training=True)
        loss = tf.reduce_mean(
            tf.keras.losses.sparse_categorical_crossentropy(
                labels, logits, from_logits=True
            )
        )
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

The tape records the forward computation. gradient computes derivatives; apply_gradients changes the weights and the optimizer's state. These are distinct operations.

training=True also enables training behaviour in layers such as dropout. A validation pass should evaluate predictions without applying updates and use the corresponding evaluation behaviour.

A successful step only improves the chosen training objective locally. Data coverage and held-out evaluation determine whether the learning transfers to useful new inputs.