PyTorch Cheatsheet

Loss Functions

Use this PyTorch reference while you build software engineering projects, review code, or refresh the syntax you reach for most.

Quick Reference Table

LossClassTask
Mean Squared Errornn.MSELossregression
Mean Absolute Errornn.L1Lossregression (robust to outliers)
Huber / Smooth L1nn.HuberLoss / nn.SmoothL1Lossregression (mixed L1/L2)
Cross-Entropynn.CrossEntropyLossmulti-class classification
Binary Cross-Entropynn.BCELossbinary / multi-label
BCE with Logitsnn.BCEWithLogitsLossbinary (numerically stable)
Negative Log-Likelihoodnn.NLLLossafter LogSoftmax
KL Divergencenn.KLDivLossdistributions, VAEs
Cosine Embeddingnn.CosineEmbeddingLosssimilarity learning
Margin Rankingnn.MarginRankingLossranking
Triplet Marginnn.TripletMarginLossmetric learning
Hinge Embeddingnn.HingeEmbeddingLossSVM-style
Connectionist Temporalnn.CTCLosssequence-to-sequence (ASR)
Poisson NLLnn.PoissonNLLLosscount data
Gaussian NLLnn.GaussianNLLLosslearned variance regression
Multi-Label Marginnn.MultiLabelMarginLossmulti-label ranking
Soft Marginnn.SoftMarginLossbinary SVM-style

Regression Losses

import torch.nn as nn

# MSE — L2 loss
criterion = nn.MSELoss(reduction='mean')   # 'mean' | 'sum' | 'none'
loss = criterion(preds, targets)           # both shape (N, ...) float

# MAE — L1 loss
criterion = nn.L1Loss(reduction='mean')

# Huber loss — L2 near 0, L1 far from 0
criterion = nn.HuberLoss(reduction='mean', delta=1.0)

# Smooth L1 (PyTorch's original Huber variant, beta=1 by default)
criterion = nn.SmoothL1Loss(reduction='mean', beta=1.0)
# Note: SmoothL1Loss(beta=δ) == HuberLoss(delta=δ) / δ (scaled)

reduction='none' returns per-element losses — useful for sample weighting:

per_sample = criterion(preds, targets)         # shape (N,) or (N, ...)
weighted = (per_sample * sample_weights).mean()

Classification Losses

Cross-Entropy (multi-class)

# Raw logits in, one integer label per sample
criterion = nn.CrossEntropyLoss(
    weight=None,          # class weights tensor, shape (C,)
    ignore_index=-100,    # ignore this label index
    reduction='mean',
    label_smoothing=0.1,  # label smoothing (0.0 = off)
)
# input:  (N, C) logits  — do NOT softmax first
# target: (N,) int64 class indices  OR  (N, C) soft labels
loss = criterion(logits, targets)

# Internally does: log_softmax(logits) → nll_loss

Binary Cross-Entropy

# BCELoss — requires sigmoid applied first
sigmoid = nn.Sigmoid()
criterion = nn.BCELoss(weight=None, reduction='mean')
loss = criterion(sigmoid(logits), targets.float())

# BCEWithLogitsLoss — numerically stable, preferred
criterion = nn.BCEWithLogitsLoss(
    weight=None,
    reduction='mean',
    pos_weight=None,   # tensor of shape (C,) — upweight positive class
)
loss = criterion(logits, targets.float())
# input:  (N, *) raw logits
# target: (N, *) float {0.0, 1.0}

# pos_weight > 1 → recall focus; < 1 → precision focus
pos_weight = torch.tensor([3.0])   # positive class 3× more important
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)

NLL Loss (after LogSoftmax)

log_probs = nn.LogSoftmax(dim=1)(logits)   # (N, C)
criterion = nn.NLLLoss(weight=None, ignore_index=-100, reduction='mean')
loss = criterion(log_probs, targets)        # targets: (N,) int64

Metric Learning Losses

# Cosine Embedding Loss — y=1: embeddings should be similar, y=-1: dissimilar
criterion = nn.CosineEmbeddingLoss(margin=0.0, reduction='mean')
loss = criterion(x1, x2, y)   # x1,x2: (N, D); y: (N,) of {1,-1}

# Triplet Margin Loss — anchor, positive, negative
criterion = nn.TripletMarginLoss(margin=1.0, p=2.0, eps=1e-6, swap=False)
loss = criterion(anchor, positive, negative)   # each (N, D)

# With custom distance
criterion = nn.TripletMarginWithDistanceLoss(
    distance_function=nn.PairwiseDistance(p=2),
    margin=1.0,
    swap=False,
)

# Margin Ranking Loss — score1 should be larger than score2 by margin if y=1
criterion = nn.MarginRankingLoss(margin=0.0, reduction='mean')
loss = criterion(score1, score2, y)   # y: (N,) of {1,-1}

Sequence Losses

# CTC Loss — for variable-length output (ASR, OCR)
criterion = nn.CTCLoss(blank=0, reduction='mean', zero_infinity=True)
# log_probs: (T, N, C) — time-major
# targets:   (N, S) or (sum_of_target_lengths,)
# input_lengths: (N,) — actual input lengths
# target_lengths: (N,) — actual target lengths
loss = criterion(log_probs, targets, input_lengths, target_lengths)

Probabilistic Losses

# KL Divergence: KL(input || target) — input should be log-probabilities
criterion = nn.KLDivLoss(reduction='batchmean', log_target=False)
# input: log-probabilities (N, *)
# target: probabilities if log_target=False; log-probs if True
loss = criterion(F.log_softmax(logits, dim=1), target_probs)

# Gaussian NLL — learn output variance
criterion = nn.GaussianNLLLoss(full=False, eps=1e-6, reduction='mean')
loss = criterion(mu, targets, var)   # var: predicted variance (positive)

# Poisson NLL — for count data (log-space inputs by default)
criterion = nn.PoissonNLLLoss(log_input=True, full=False, reduction='mean')
loss = criterion(log_predictions, targets)

Functional API

All losses are also available as functions in torch.nn.functional:

import torch.nn.functional as F

F.mse_loss(input, target, reduction='mean')
F.l1_loss(input, target)
F.huber_loss(input, target, delta=1.0)
F.cross_entropy(input, target, weight=None, ignore_index=-100, label_smoothing=0.0)
F.binary_cross_entropy(input, target)
F.binary_cross_entropy_with_logits(input, target, pos_weight=None)
F.nll_loss(input, target)
F.kl_div(input, target, reduction='batchmean', log_target=False)
F.cosine_embedding_loss(x1, x2, y, margin=0.0)
F.triplet_margin_loss(anchor, positive, negative, margin=1.0)
F.ctc_loss(log_probs, targets, input_lengths, target_lengths)

Custom Loss Functions

# Simple functional approach
def focal_loss(logits, targets, gamma=2.0, alpha=0.25):
    bce = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')
    pt = torch.exp(-bce)
    focal = alpha * (1 - pt) ** gamma * bce
    return focal.mean()

# nn.Module approach (needed if the loss has learnable parameters)
class DiceLoss(nn.Module):
    def __init__(self, smooth=1.0):
        super().__init__()
        self.smooth = smooth

    def forward(self, logits, targets):
        probs = torch.sigmoid(logits)
        intersection = (probs * targets).sum(dim=(2, 3))
        denom = probs.sum(dim=(2, 3)) + targets.sum(dim=(2, 3))
        dice = (2 * intersection + self.smooth) / (denom + self.smooth)
        return 1 - dice.mean()

Class Imbalance Strategies

# 1. Class weights in CrossEntropyLoss
counts = torch.bincount(train_labels)
weights = 1.0 / counts.float()
weights = weights / weights.sum()
criterion = nn.CrossEntropyLoss(weight=weights.to(device))

# 2. Positive weighting in BCEWithLogitsLoss
pos_weight = (neg_count / pos_count) * torch.ones(1)
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)

# 3. Label smoothing
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

# 4. Focal loss (see custom example above)

Gotchas

nn.CrossEntropyLoss expects raw logits — do not apply softmax first. It fuses log_softmax + nll_loss internally for numerical stability.

nn.BCELoss requires inputs in [0, 1] — apply sigmoid first. Prefer nn.BCEWithLogitsLoss which does this internally and is more numerically stable.

nn.KLDivLoss with reduction='mean' divides by total elements (N*C), not batch size. Use reduction='batchmean' for mathematically correct KL divergence.

NLLLoss expects log-probabilities as input, not raw probabilities. Pair it with LogSoftmax.