5.2.6 Loss Functions

Why does a network need a loss function?

In Activation Functions, we left the last layer of a network without an activation and said the loss would take care of it. This page explains that promise, and finishes the picture of how a network learns.

A network turns an input into some output numbers. To improve, it needs one number that says how wrong those outputs are over the training examples. That number is the loss, and it is exactly what Gradient Descent makes smaller, step by step, using the gradients that Backpropagation computes. Choose the wrong loss and the network carefully learns the wrong thing.

You already know the ideas. Loss Functions in the Classical Machine Learning part covered what a loss is, the difference between a loss and a metric, and the main losses: squared error, absolute error, Huber loss, and binary and multiclass cross-entropy. We will not repeat that theory here. This page is about using those losses in PyTorch: which class to pick for each kind of task, what shapes and types it expects, and the mistakes that make training fail without any error message.

What you need to know

By the end of this module, you should be comfortable with these ideas:

That is enough to train every model in the following modules.

Learn it properly

This page shows the losses in action. The readings explain where cross-entropy comes from, which will make its behaviour feel natural rather than magic.

Pick one main course

Our pick: Dive into Deep Learning, section 4.1.2

Read section 4.1.2, Loss Function, in Dive into Deep Learning. It shows how cross-entropy follows from asking the model to give high probability to the right answers, and how the softmax and the loss fit together so neatly that their gradient is simply "predicted probability minus the true one". Section 4.1.1 just before it introduces softmax, if you want the full story.

Section 4.1.3, on information theory, is interesting but optional; you do not need it for anything in this guide.

If you prefer something else

• Videos: StatQuest's ArgMax and SoftMax followed by Cross Entropy. They are short and build the idea slowly.
• A free book: chapter 3 of Michael Nielsen's Neural Networks and Deep Learning, "The cross-entropy cost function", explains why cross-entropy trains classifiers faster than squared error.
• Course notes: Stanford CS231n's Linear Classification notes, the Softmax classifier part.
• Reference: the loss functions section of the PyTorch docs. Each loss's page lists the exact shapes it expects; check it whenever you use a loss for the first time.

A quick tour

The rest of this page goes through the losses by kind of task, with small tensors you can check by hand. For each one, pay as much attention to the shape and dtype of the predictions and targets as to the loss itself; that is where most bugs come from.

Logits, shapes and dtypes

Our networks end with a plain nn.Linear layer, so they output raw scores called logits. A batch of $N$ examples gives:

Keep this table open while you write your first few training loops. We will now go through it row by row. Every loss below is a class: we create it once, then call it like a function with (predictions, targets).

Regression

For predicting numbers we use the losses from the classical part: squared error, absolute error and Huber loss.

import torch
from torch import nn

predictions = torch.tensor([2.5, 0.0, 2.0, 8.0])
targets = torch.tensor([3.0, -0.5, 2.0, 7.0])

mse = nn.MSELoss()
mae = nn.L1Loss()
huber = nn.HuberLoss(delta=1.0)

print(mse(predictions, targets))
print(mae(predictions, targets))
print(huber(predictions, targets))

tensor(0.3750)
tensor(0.5000)
tensor(0.1875)

Let us check the first one by hand. The errors are −0.5, 0.5, 0 and 1; their squares are 0.25, 0.25, 0 and 1; and the average is 1.5 / 4 = 0.375. The classical module explains when to prefer each one; nn.MSELoss is the usual starting point.

Binary classification

For a yes/no task, the network outputs one logit per example, and the targets are 1.0 for "yes" and 0.0 for "no". The loss is binary cross-entropy, computed straight from the logits:

logits = torch.tensor([2.0, -1.0, 0.5])
labels = torch.tensor([1.0, 0.0, 0.0])

binary_loss = nn.BCEWithLogitsLoss()
print(binary_loss(logits, labels))

tensor(0.4714)

"With logits" means the loss applies the sigmoid itself. We can check that by doing the two steps separately, with nn.BCELoss, which expects probabilities instead:

probabilities = torch.sigmoid(logits)
print(probabilities)
print(nn.BCELoss()(probabilities, labels))

tensor([0.8808, 0.2689, 0.6225])
tensor(0.4714)

Same answer. So why not always use sigmoid and BCELoss? Because the two-step version breaks for very confident predictions. Take one example with logit 200 whose true label is 0; the correct loss is about 200:

confident_logit = torch.tensor([200.0])
true_label = torch.tensor([0.0])

print(nn.BCEWithLogitsLoss()(confident_logit, true_label))
print(torch.sigmoid(confident_logit))
print(nn.BCELoss()(torch.sigmoid(confident_logit), true_label))

tensor(200.)
tensor([1.])
tensor(100.)

The sigmoid of 200 rounds to exactly 1.0 in 32-bit floats, and $\log(1 - 1)$ is minus infinity, so BCELoss has to cut it off at 100. The combined loss never forms that probability and gets the right answer. Always output logits and use nn.BCEWithLogitsLoss; apply torch.sigmoid only when you want probabilities to look at or submit.

Multi-label classification

Sometimes each example can have several labels at once: a photo can contain both a cat and a dog. That is just $K$ binary problems side by side. The model outputs $K$ logits per example, the targets are a float grid of 0s and 1s with the same shape (N, K), and the loss is the same nn.BCEWithLogitsLoss, averaged over every entry. Nothing else changes.

Multiclass classification

When each example belongs to exactly one of $K$ classes, the network outputs $K$ logits per example, and the loss is multiclass cross-entropy. The targets are not one-hot rows: they are the class numbers themselves, as integers.

Let us take two examples and three classes. The first example belongs to class 0 and the second to class 2:

scores = torch.tensor([[2.0, 0.5, -1.0],
                       [0.1, 0.2, 3.0]])
true_classes = torch.tensor([0, 2])

multiclass_loss = nn.CrossEntropyLoss()
print(multiclass_loss(scores, true_classes))

tensor(0.1755)

nn.CrossEntropyLoss does two things in one: it turns each row of logits into probabilities with softmax (in a numerically safe way), and then takes minus the log of the probability given to the true class. We can do the two steps ourselves to see this:

import torch.nn.functional as F

log_probabilities = F.log_softmax(scores, dim=1)
print(log_probabilities)

tensor([[-0.2413, -1.7413, -3.2413],
        [-3.0096, -2.9096, -0.1096]])

The true classes are 0 and 2, so we pick out −0.2413 from the first row and −0.1096 from the second. Minus their average is $(0.2413 + 0.1096)/2 \approx 0.1755$, the same number as before. dim=1 means "softmax across each row", that is, across the classes of one example.

To turn logits into probabilities for your own use, apply softmax yourself:

torch.softmax(scores, dim=1)

tensor([[0.7856, 0.1753, 0.0391],
        [0.0493, 0.0545, 0.8962]])

Each row adds up to 1, and the largest value in each row is the predicted class. For a class prediction alone, scores.argmax(dim=1) gives the same answer without softmax, because softmax never changes which score is largest.

Averaging, summing, or keeping each loss

By default every loss returns the average over the batch. The reduction option changes that: "sum" adds the losses up, and "none" keeps one loss per example, which is handy for finding the examples the model gets most wrong:

per_example = nn.CrossEntropyLoss(reduction="none")
print(per_example(scores, true_classes))

tensor([0.2413, 0.1096])

Leave it at the default "mean" for training. With "sum", the size of the loss (and so of the gradient steps) grows with the batch size, which makes the learning rate harder to choose.

Rare classes

If one class is much rarer than the others, a network can get a low loss by mostly ignoring it. The fix is the same idea as class weights in the classical part: make mistakes on the rare class cost more.

• nn.CrossEntropyLoss(weight=...) takes one weight per class. A common choice is to make each weight inversely proportional to how often that class appears.
• nn.BCEWithLogitsLoss(pos_weight=...) multiplies the loss on the positive examples. A common choice is the number of negative examples divided by the number of positive ones: with 900 negatives and 100 positives, pos_weight=torch.tensor([9.0]).

class_weights = torch.tensor([1.0, 1.0, 5.0])   # class 2 is rare
weighted_loss = nn.CrossEntropyLoss(weight=class_weights)
print(weighted_loss(scores, true_classes))

tensor(0.1316)

The number changed because the second example, which belongs to the rare class 2, now counts five times as much in the weighted average. Weights change what the model is trained to care about, so always check the effect on the metric that is actually scored; Model Evaluation Metrics covers the metrics used for imbalanced data.

Label smoothing

nn.CrossEntropyLoss(label_smoothing=0.1) tells the network that the true class is only 90% certain, spreading the remaining 10% over all the classes. It stops the network from becoming overconfident and sometimes improves accuracy a little. You do not need it now; it is one of the tricks worth trying once a model already works.

Watching the loss during training

In a training loop, the loss is the one number you compute every step. It is also the first thing to look at when something goes wrong. Record the average training loss and the loss on held-out validation data after every epoch, then plot both curves, as in the Matplotlib module:

• Both curves falling: training is working; keep going.
• Training loss falling, validation loss rising: the model is starting to memorise the training data; see Underfitting and Overfitting.
• Neither falling: the learning rate may be wrong, or there may be a bug; the next section lists the usual suspects.
• The loss becomes nan: the numbers have blown up, often because the learning rate is far too high.

Remember that the loss is what the network optimises, not necessarily what the task scores. A model with lower cross-entropy usually has better accuracy, but not always; the classical module explains the difference between a loss and an evaluation metric.

Common mistakes

Losses produce some of PyTorch's least friendly error messages, so it is worth knowing them in advance. Each message below is exactly what PyTorch prints.

When a loss complains, print predictions.shape, predictions.dtype, targets.shape, targets.dtype and compare them with the table in Logits, shapes and dtypes. That settles most problems in a minute.

How to practise

• Read the Dive into Deep Learning section and try its exercises on cross-entropy; the proofs are optional.
• Recompute every loss on this page by hand, or with a few lines of plain Python, until the numbers match. Cross-entropy stops feeling mysterious once you have done it twice.
• Break things on purpose: give nn.CrossEntropyLoss float targets, or a model output with the wrong shape, and read the error. Next time you meet it for real, you will recognise it.
• In A Simple MLP in PyTorch, plot the training and validation loss of a real model and describe what the curves say.

Check yourself

What comes next

That completes Deep Learning Fundamentals: a network (Multi-Layer Perceptrons), a way to improve it (Gradient Descent and Backpropagation), the maths it runs on (Tensors, Matrices, and Related Mathematical Functions), its non-linearities (Activation Functions) and what it minimises (this page).

• Dataset and DataLoader: feeding real data to a network in batches, the first step of the full PyTorch pipeline.