PyTorch Loss Functions and Optimizers
PyTorch loss functions measure how wrong a model's predictions are, and optimizers use the gradients of that loss to update the weights. Picking the right loss for your task, and a sensible optimizer and learning rate, is most of what "configuring training" means.
Regression: MSELoss
import torch
from torch import nn
loss_fn = nn.MSELoss()
pred = torch.tensor([2.5, 0.0, 2.0])
target = torch.tensor([3.0, -0.5, 2.0])
print(loss_fn(pred, target))tensor(0.1667)
Mean squared error averages the squared differences: (0.25 + 0.25 + 0) / 3 = 0.1667. Use it, or nn.L1Loss, when predicting continuous numbers such as prices.
Classification: CrossEntropyLoss
For multi-class problems, pass raw scores (logits) of shape [batch, classes] and integer labels of shape [batch]. CrossEntropyLoss applies log-softmax internally, so do not add a softmax layer yourself.
loss_fn = nn.CrossEntropyLoss()
logits = torch.tensor([[2.0, 0.5, -1.0],
[0.1, 0.2, 3.0]])
labels = torch.tensor([0, 2])
print(loss_fn(logits, labels))tensor(0.1755)
Binary Classification: BCEWithLogitsLoss
loss_fn = nn.BCEWithLogitsLoss()
logits = torch.tensor([1.5, -2.0])
labels = torch.tensor([1.0, 0.0])
print(loss_fn(logits, labels))tensor(0.1641)
Labels are floats here. BCEWithLogitsLoss combines a sigmoid and binary cross-entropy in a numerically stable way, which is better than nn.Sigmoid followed by nn.BCELoss.
Optimizers
model = nn.Linear(10, 1)
opt_sgd = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
opt_adam = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
# one update step
x, y = torch.randn(8, 10), torch.randn(8, 1)
loss = nn.MSELoss()(model(x), y)
opt_adam.zero_grad()
loss.backward()
opt_adam.step()
print("updated, loss was", round(loss.item(), 3))updated, loss was 1.284 (sample output — random data)
- SGD with momentum: simple, strong for CNNs when tuned, needs a higher lr (0.01-0.1).
- Adam: adapts the step per parameter, works well out of the box (lr around 1e-3).
- AdamW: Adam with correct weight decay, the default choice for transformers and most modern models.
Learning Rate Schedulers
Lowering the learning rate during training often improves final accuracy. A scheduler adjusts it for you; call scheduler.step() once per epoch (or per batch, depending on the scheduler).
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10)
for epoch in range(3):
# ... train one epoch ...
sched.step()
print(epoch, round(sched.get_last_lr()[0], 6))0 0.000976 1 0.000905 2 0.000794
Frequently asked questions
Which loss function should I use in PyTorch?
MSELoss or L1Loss for regression, CrossEntropyLoss for single-label multi-class classification, and BCEWithLogitsLoss for binary or multi-label classification.
Should I apply softmax before CrossEntropyLoss?
No. CrossEntropyLoss expects raw logits and applies log-softmax itself. Adding softmax first makes training worse.
Adam or SGD: which optimizer is better?
Adam and AdamW converge quickly with little tuning and are the usual default. Well-tuned SGD with momentum can generalize slightly better for some vision models.