PyTorch Loss Functions

A loss function measures how wrong a model's predictions are. During training, the optimizer reads the loss value and adjusts the model's weights to make it smaller. Picking the wrong loss function for your task causes the model to optimize for the wrong goal — even if every other part of the pipeline is correct.

Loss Function Workflow

Model Output (predictions)
         │
         ▼
   Loss Function  ← compares predictions to true labels
         │
         ▼
   Scalar Loss Value  ← a single number measuring total error
         │
         ▼
   loss.backward()  ← computes gradients
         │
         ▼
   optimizer.step()  ← weights updated to reduce loss

Mean Squared Error – For Regression

MSE computes the average of squared differences between predictions and true values. Use it when predicting continuous numbers — house prices, temperature, stock values.

import torch
import torch.nn as nn

loss_fn = nn.MSELoss()

predictions = torch.tensor([2.5, 3.0, 4.1])
targets     = torch.tensor([3.0, 3.0, 4.0])

loss = loss_fn(predictions, targets)
print(loss)   # tensor(0.0867)

# Manual calculation:
# ((2.5-3)² + (3.0-3)² + (4.1-4.0)²) / 3
# = (0.25 + 0 + 0.01) / 3 = 0.0867

Binary Cross-Entropy – For Binary Classification

Binary Cross-Entropy (BCE) measures the difference between two probability distributions when there are two classes (yes/no, cat/not-cat, spam/not-spam). Combine it with sigmoid output.

loss_fn = nn.BCELoss()

# Predictions must be probabilities (0 to 1) from sigmoid
pred = torch.tensor([0.9, 0.2, 0.8, 0.1])   # predicted probabilities
true = torch.tensor([1.0, 0.0, 1.0, 0.0])   # true labels: 1=positive, 0=negative

loss = loss_fn(pred, true)
print(loss)   # low loss: model is mostly correct

BCEWithLogitsLoss – Preferred for Binary Classification

BCEWithLogitsLoss combines the sigmoid activation and BCE loss in one operation. This is numerically more stable than applying sigmoid first and then BCE separately. Use this instead of BCELoss in practice.

loss_fn = nn.BCEWithLogitsLoss()

raw_scores = torch.tensor([2.1, -1.5, 3.0, -0.8])   # raw logits (no sigmoid)
true_labels = torch.tensor([1.0, 0.0, 1.0, 0.0])

loss = loss_fn(raw_scores, true_labels)
print(loss)

Cross-Entropy Loss – For Multi-Class Classification

CrossEntropyLoss handles classification with three or more classes. It internally applies softmax to the model's raw output scores and then computes the negative log probability of the correct class.

loss_fn = nn.CrossEntropyLoss()

# Model outputs raw scores for 3 classes (batch of 2 samples)
logits = torch.tensor([[2.0, 0.5, 0.1],    # sample 1: class 0 most likely
                        [0.1, 0.2, 3.0]])   # sample 2: class 2 most likely

labels = torch.tensor([0, 2])   # true class indices

loss = loss_fn(logits, labels)
print(loss)   # low value since predictions match labels

CrossEntropyLoss diagram:
  Raw scores: [2.0, 0.5, 0.1]  for sample 1
      │
      ▼ softmax (internal)
  Probabilities: [0.76, 0.17, 0.07]
      │
      ▼ negative log of correct class (class 0 = 0.76)
  Loss: -log(0.76) ≈ 0.27   ← small: model was confident and correct

Negative Log Likelihood Loss

NLLLoss is used when your model outputs log-probabilities (after applying log_softmax). It is equivalent to cross-entropy but requires you to apply log_softmax yourself first.

import torch.nn.functional as F

log_probs = F.log_softmax(logits, dim=1)
loss_fn = nn.NLLLoss()
loss = loss_fn(log_probs, labels)
print(loss)

Choosing the Right Loss Function

Task                              Loss Function
───────────────────────────────────────────────────────
Predict a number (regression)     nn.MSELoss
Predict yes/no (binary)           nn.BCEWithLogitsLoss
Predict one of N classes          nn.CrossEntropyLoss
Custom / flexible                 write your own

Custom Loss Function

You can write your own loss function using standard tensor operations. Any function that returns a scalar can serve as a loss.

def huber_loss(pred, target, delta=1.0):
    error = pred - target
    abs_error = error.abs()
    quadratic = torch.clamp(abs_error, max=delta)
    linear = abs_error - quadratic
    return (0.5 * quadratic ** 2 + delta * linear).mean()

loss = huber_loss(predictions, targets)
loss.backward()

Summary

Loss functions measure prediction error. MSELoss suits regression tasks. BCEWithLogitsLoss handles binary classification with raw logits. CrossEntropyLoss handles multi-class classification. Choose based on your output type: continuous values need MSE, binary classification needs BCE, multi-class classification needs CrossEntropy. Always call loss.backward() after computing the loss to get gradients. You can write custom loss functions as regular Python functions that return a scalar tensor.