Supplement · Loss Functions

Loss Functions in PyTorch

Colab Notebook · ~45 min
Google Colab Notebook
Loss Functions in PyTorch
Python · ~45 min
Open in Colab
Lab Objectives
1
Implement each of the 20 PyTorch loss functions from scratch using tensor operations, then verify against nn.* equivalents
2
Derive and verify the closed-form gradient of L1, MSE, and BCE losses; compare against loss.backward()
3
Build a toy training loop for each loss family and observe convergence behaviour
4
Implement a numerically stable cross-entropy using the log-sum-exp trick and verify against nn.CrossEntropyLoss
5
Train a Siamese network with TripletMarginLoss on MNIST embeddings and visualise the learned metric space

Lab Overview

This notebook ties every formula from the readings to runnable code. For each loss you will:

  1. Write a from-scratch implementation using basic PyTorch tensor ops (torch.abs, torch.exp, torch.log, etc.)
  2. Compare numerically against the corresponding torch.nn loss
  3. Verify the gradient by running .backward() and comparing to the analytical gradient

By the end you will have built intuition for why PyTorch's built-in losses use the specific formulations they do — especially the numerically stable variants.

Sections

Section Losses covered
1 L1Loss, MSELoss, HuberLoss, SmoothL1Loss
2 PoissonNLLLoss, GaussianNLLLoss
3 BCELoss, BCEWithLogitsLoss (log-sum-exp derivation)
4 CrossEntropyLoss, NLLLoss (log-softmax + NLL)
5 SoftMarginLoss, MultiLabelSoftMarginLoss
6 MultiMarginLoss, MultiLabelMarginLoss
7 KLDivLoss (batchmean vs mean reduction)
8 MarginRankingLoss, HingeEmbeddingLoss, CosineEmbeddingLoss
9 TripletMarginLoss, TripletMarginWithDistanceLoss
10 End-to-end: Siamese MNIST with triplet loss