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
| Loss | Class | Task |
|---|---|---|
| Mean Squared Error | nn.MSELoss | regression |
| Mean Absolute Error | nn.L1Loss | regression (robust to outliers) |
| Huber / Smooth L1 | nn.HuberLoss / nn.SmoothL1Loss | regression (mixed L1/L2) |
| Cross-Entropy | nn.CrossEntropyLoss | multi-class classification |
| Binary Cross-Entropy | nn.BCELoss | binary / multi-label |
| BCE with Logits | nn.BCEWithLogitsLoss | binary (numerically stable) |
| Negative Log-Likelihood | nn.NLLLoss | after LogSoftmax |
| KL Divergence | nn.KLDivLoss | distributions, VAEs |
| Cosine Embedding | nn.CosineEmbeddingLoss | similarity learning |
| Margin Ranking | nn.MarginRankingLoss | ranking |
| Triplet Margin | nn.TripletMarginLoss | metric learning |
| Hinge Embedding | nn.HingeEmbeddingLoss | SVM-style |
| Connectionist Temporal | nn.CTCLoss | sequence-to-sequence (ASR) |
| Poisson NLL | nn.PoissonNLLLoss | count data |
| Gaussian NLL | nn.GaussianNLLLoss | learned variance regression |
| Multi-Label Margin | nn.MultiLabelMarginLoss | multi-label ranking |
| Soft Margin | nn.SoftMarginLoss | binary 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.CrossEntropyLossexpects raw logits — do not applysoftmaxfirst. It fuseslog_softmax+nll_lossinternally for numerical stability.
nn.BCELossrequires inputs in[0, 1]— applysigmoidfirst. Prefernn.BCEWithLogitsLosswhich does this internally and is more numerically stable.
nn.KLDivLosswithreduction='mean'divides by total elements (N*C), not batch size. Usereduction='batchmean'for mathematically correct KL divergence.
NLLLossexpects log-probabilities as input, not raw probabilities. Pair it withLogSoftmax.