🎓 A1Learner
🔍

PyTorch Cross-Entropy Loss

Cross-entropy loss is the loss for picking one correct class from several. It takes logits, the raw scores a model gives before any probabilities, plus the number of the correct class.

Scores and the correct class

The row holds one score per class. The target holds the position of the correct class.

import torch

scores = torch.tensor([[3.0, 0.0, 0.0]])
target = torch.tensor([0])
print(scores)
print(target)

A confident right answer

The biggest score sits on the correct class here. The loss stays small.

import torch

scores = torch.tensor([[3.0, 0.0, 0.0]])
target = torch.tensor([0])
loss_fn = torch.nn.CrossEntropyLoss()
print(loss_fn(scores, target))

A confident wrong answer

The biggest score now sits on a wrong class. The loss jumps up.

import torch

scores = torch.tensor([[0.0, 0.0, 3.0]])
target = torch.tensor([0])
loss_fn = torch.nn.CrossEntropyLoss()
print(loss_fn(scores, target))

Example

import torch

loss_fn = torch.nn.CrossEntropyLoss()

right_scores = torch.tensor([[3.0, 0.0, 0.0]])
right_target = torch.tensor([0])
right_loss = loss_fn(right_scores, right_target)

wrong_scores = torch.tensor([[0.0, 0.0, 3.0]])
wrong_target = torch.tensor([0])
wrong_loss = loss_fn(wrong_scores, wrong_target)

print(f"Confident and right: {right_loss.item():.4f}")
print(f"Confident and wrong: {wrong_loss.item():.4f}")

Key points

  • Class: one of the fixed choices the model picks from, like cat or dog
  • Logits: the raw scores a model gives before any probabilities
  • Target: the number of the correct class, starting at 0
  • Use nn.CrossEntropyLoss() with a row of scores and one target per row
  • Cross-entropy applies its own softmax step inside, so no softmax goes before it
  • A confident wrong answer earns a much bigger loss than a confident right one

In this example, one loss class scores two predictions. The first row gives the correct class the biggest score by far. Its loss is small, at 0.0949. The second row gives a wrong class the biggest score instead. Its loss is much bigger, at 3.0949. Both rows use raw scores, with no softmax before the loss.

Note

The scores are floats, and the target is a whole number, not a float

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does cross-entropy loss need as the target for each row?