7  Training and inference

7.1 Training loss

To train our classifier, we need images together with their correct class labels. For MNIST, each training image is paired with a digit label from 0 to 9. The model predicts a probability for each possible digit, and the loss function compares this prediction with the supplied label.

In ENGSCI 205, you used the negative log-likelihood loss. For one image with correct class label \(y\), this is

\[ \ell=-\log p_y, \]

where \(p_y\) is the probability the model assigns to the observed class \(y\) for this image, and \(\log\) denotes the natural logarithm.

For example, suppose an image is labelled as digit 7:

Probability assigned to digit 7 Loss
\(p_7=0.10\) \(-\log(0.10)\approx2.303\)
\(p_7=0.80\) \(-\log(0.80)\approx0.223\)

We can interpret the loss as how surprising the correct label is according to the model. If the model assigns probability 1 to the correct class, the loss is zero. As that probability approaches zero, the loss becomes arbitrarily large.

We can score predictions in different ways. For example, the simplest way is to assign 1 if the highest-probability class matches the supplied label, and 0 otherwise. This leads to the classification accuracy, which is the average of these 0/1 scores, i.e., the fraction of highest-probability predictions that are correct.

As an alternative to a 0/1 score, negative log-likelihood distinguishes between predictions that select the same class but assign different probabilities to the correct class. For example, two predictions might both correctly select digit 7 as the highest probability class, while assigning it probabilities 0.9 and 0.6. Both count as correct, but the first has a lower negative log-likelihood loss.

The logarithm connects this loss to maximum likelihood. Treating the labelled examples as independent, we multiply the probabilities assigned to their observed labels. Taking logarithms turns this product into a sum. For two images whose correct labels receive probabilities \(p_a\) and \(p_b\),

\[ -\log(p_a p_b)=-\log p_a-\log p_b. \]

Maximising the likelihood is therefore equivalent to minimising the sum of these negative log-probabilities. For a minibatch, we use their mean, giving the loss per image.

7.2 Cross-entropy and one-hot targets

Suppose an image has the correct digit label 7. So far, we have represented its target using that single number. Another representation uses a vector with one entry for each possible digit, ordered from 0 to 9:

\[ q=(0,0,0,0,0,0,0,1,0,0). \]

The entry corresponding to digit 7 is 1, and all other entries are zero. This is called a one-hot vector, because exactly one entry is “on”. We can interpret \(q\) as a target probability distribution with all its probability mass at the supplied class label. Its entries are non-negative and sum to one.

The target distribution \(q\) represents the observed label. The model supplies a separate predicted distribution \(p\), which we compare with \(q\) using the loss function.

The categorical cross-entropy loss is

\[ \ell=-\sum_{c=0}^{9}q_c\log p_c. \]

For each class, we multiply the model’s log-probability by the corresponding target entry, then sum and negate the result. Since all target entries except \(q_7\) are zero, only the digit-7 term remains:

\[ \ell=-q_7\log p_7=-\log p_7. \]

Thus, categorical cross-entropy with a one-hot target gives exactly the negative log-likelihood loss introduced above. It selects the probability assigned to the observed class.

7.3 LogSoftmax and NLLLoss

PyTorch’s nn.NLLLoss expects log-probabilities, rather than probabilities or raw class scores. We therefore use nn.LogSoftmax(dim=1) after the final linear layer. It combines the softmax calculation and logarithm into one operation.

NLLLoss takes these log-probabilities together with the correct class labels. For each image, it selects the log-probability of the correct class and negates it, then averages the losses across the minibatch.

In our PyTorch code, we supply the digit labels directly; we do not need to construct the one-hot vectors ourselves. The labels are supplied only to the loss calculation. The network must predict its class scores from the images themselves.

7.4 Updating the parameters

This follows the training procedure from the ENGSCI 205 course book and associated ENGSCI 205 Lab 5. ENGSCI 205 resources are available on Canvas.

Our earlier example separated the convolutional block, model, from the final classifier. For training, we combine both with LogSoftmax:

network = nn.Sequential(
    model,
    classifier,
    nn.LogSoftmax(dim=1),
)

criterion = nn.NLLLoss()
optimizer = torch.optim.SGD(network.parameters(), lr=0.01)

Here network.parameters() includes the parameters of both the convolutional layer and the final fully connected layer. The optimiser updates these parameters; lr specifies the learning rate.

For the following code fragment, suppose we have loaded a minibatch of labelled digit images:

  • mb_data has shape (batch, 1, 28, 28);
  • mb_targets has shape (batch,), containing the correct digit label for each image.

One training update consists of:

optimizer.zero_grad()

log_probabilities = network(mb_data)
mb_loss = criterion(log_probabilities, mb_targets)

mb_loss.backward()
optimizer.step()

These steps do different jobs:

  1. zero_grad() clears gradients from the previous update. It does not reset the model parameters.
  2. network(mb_data) computes predictions using the current parameters—the forward pass.
  3. The loss calculation compares those predictions with the supplied labels and returns the mean loss over the minibatch.
  4. backward() computes the gradient of the loss with respect to the parameters, using backpropagation.
  5. step() changes the parameters using those gradients.

Computing predictions and evaluating their loss do not themselves change the parameters. The update happens when we call optimizer.step().

7.5 Minibatches and epochs

Rather than use the entire training dataset for every parameter update, we usually process smaller groups of examples called minibatches. For each minibatch, we calculate the mean loss, compute its gradients and update the parameters.

An epoch is one pass through the training dataset. The parameters are updated after each minibatch, so one epoch usually contains many updates.

For example, MNIST has 60,000 training images. With minibatches of 32 images, one epoch contains

\[ \frac{60{,}000}{32}=1875 \]

parameter updates. Training for ten epochs therefore involves 18,750 updates, with each image used once per epoch.

In PyTorch, a DataLoader supplies the minibatches. Iterating over a training loader returns a pair of tensors, mb_data and mb_targets, at each step. If the loader was created with shuffle=True, it shuffles the training examples before each new pass.

Once we have constructed train_loader, as in the ENGSCI 205 example, we can repeat our update using two loops:

num_epochs = 10
network.train()

for epoch in range(num_epochs):
    for mb_data, mb_targets in train_loader:
        optimizer.zero_grad()

        log_probabilities = network(mb_data)
        mb_loss = criterion(log_probabilities, mb_targets)

        mb_loss.backward()
        optimizer.step()

network.train() selects training mode, reversing the effect of network.eval(). It does not itself train the network; the parameter updates still occur inside the loop.

The inner loop processes the minibatches within one epoch. The outer loop repeats this process for the chosen number of epochs, continuing from the parameters learned so far.

7.6 Inference

Once a network has been trained, we can use it to make predictions for new images. In machine learning this is, unfortunately, called inference (in statistics, inference usually refers to estimating unknown parameters and quantifying uncertainty, while prediction refers to the… prediction component). Here we use the learned parameters without updating them, and do not need class labels to produce the predictions.

For a new batch of digit images, prepared in the same way as the training images:

network.eval()

with torch.no_grad():
    log_probabilities = network(new_images)
    probabilities = log_probabilities.exp()
    predicted_digits = probabilities.argmax(dim=1)

Our network ends with LogSoftmax, so its outputs are log-probabilities. Taking their exponential recovers the probabilities. The argmax(dim=1) operation returns the index of the largest probability for each image. Because our classes are ordered from 0 to 9, these indices give the predicted digits.

The first two lines have different purposes:

  • network.eval() sets the network to evaluation mode. Some layers behave differently during training and inference. For example, dropout randomly sets some activations to zero during training, but stops doing this in evaluation mode. Our small network does not contain dropout, but other architectures may.
  • torch.no_grad() tells PyTorch not to record the operations needed for gradient calculations. We do not need these when making predictions, and omitting them reduces memory use.

Neither operation trains the network or measures its accuracy. To assess prediction quality, we would compare the predictions with known labels for images that were not used to fit the model parameters.

7.7 Pretrained models

We do not always need to train a network ourselves. A pretrained model provides parameters already learned from another dataset. Loading those parameters allows us to begin with inference.

This is what we did when loading yolo11n.pt through Ultralytics. The detector had already been trained to recognise its object categories before we supplied the basketball frame. Our call to model.predict(...) used those existing parameters; it did not train the detector on our video.

7.8 Transfer learning and fine-tuning

Transfer learning reuses what a model has learned from one dataset or task when developing a model for another. Rather than initialise every parameter randomly, we start with some parameters already learned from data.

For example, suppose we want to adapt a pretrained image classifier to a new set of categories. We can retain its convolutional layers but replace the final classification layer with one that produces scores for our new categories. We then train the new classification layer using labelled examples of those categories. If we keep the pretrained parameters fixed, the existing convolutional layers act as a feature extractor.

We can also continue training some or all of the pretrained layers. This is called fine-tuning. It uses the same forward-pass, loss, backpropagation and parameter-update steps described above, but starts from pretrained parameters rather than a random initialisation.

Fine-tuning allows the features themselves to change in response to the new training data. Whether this improves predictions should be checked on separate images that were not used to update the parameters.

In our basketball example, we have so far used YOLO’s pretrained parameters without changing them. Selecting the person category or changing the confidence threshold changes which predictions are returned, but does not train the model. Similarly, filtering detections according to their estimated court positions is a separate processing step, not fine-tuning.

Our MNIST classifier produces one vector of class scores for each image. For the basketball task, we instead need locations and class scores for multiple objects within the same image. In the next chapter, we consider how an object detector uses learned image features to predict candidate bounding boxes and class scores, and how these candidates are processed to produce the returned detections.