PyTorch Cheatsheet

Loss Functions

Use this PyTorch reference while you build software engineering projects, review code for technical interview prep, or polish examples for a software engineer resume.

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.