Deep Learning with PyTorch: First Steps

Book 11 of 50 — AstolixGen Learning Series

---

About This Book

If you are an MS or PhD student, or an early AI researcher, there is one skill that will quietly decide how fast you can turn an idea into a published result: the ability to build, train, and debug a neural network yourself, in code. Papers do not come from reading about backpropagation. They come from running experiments — dozens of them — and knowing exactly what your code is doing at every step.

This book teaches you PyTorch, the deep learning framework that has become the working language of AI research. Almost every major result you read about in the last decade — from transformers to diffusion models — was prototyped in PyTorch. Conferences like NeurIPS, ICML, and ICLR are full of PyTorch code, and when a paper releases its code, it is most often a PyTorch repository. Learning PyTorch is not about learning one tool among many; it is about joining the conversation that research actually happens in.

You do not need any prior deep learning framework experience. This book assumes you know Python and basic machine learning concepts (what a model is, what training and test data are, what a loss function means at a high level — covered in Books 2 and 7 of this series). Everything else is built from scratch, one small working program at a time.

Every chapter in this book follows the same honest structure: a clear explanation in plain language, working code you can run immediately (mostly on a free Google Colab GPU), a "For Your Research" box that connects the chapter's ideas to real publication practice, and key takeaways. By the end, you will not just "know PyTorch" — you will have the habits of a researcher who runs clean, reproducible experiments: setting seeds, logging results, saving checkpoints, and organizing code so that six months later you (or a reviewer) can reproduce everything.

A note on honesty, because research demands it: this book contains no invented facts. Code examples are real, minimal, and tested against PyTorch's documented behavior. References are real papers you can look up. When something is a simplification, the book says so.

Prerequisites: Python (Book 3 of this series), basic machine learning concepts (Book 2), neural network basics (Book 7), and gradient descent intuition (Book 8). A Google account (for free Colab GPUs) is helpful but optional.

---

Learning Objectives

By the end of this book, you will be able to:

1. Explain why PyTorch is the standard framework for AI research, and describe how it differs from TensorFlow and JAX in the ways that matter to a researcher.
2. Set up a working PyTorch environment on your own machine (CPU or GPU) and on free cloud resources like Google Colab, and verify the installation.
3. Create and manipulate tensors — PyTorch's core data structure — including creation, indexing, reshaping, broadcasting, and moving tensors between CPU and GPU.
4. Explain automatic differentiation (autograd) in plain terms and use it to compute gradients in your own code.
5. Build a neural network with nn.Module, understanding parameters, layers, activation functions, and the forward pass.
6. Choose appropriate loss functions and optimizers for common research tasks (classification, regression) and justify your choices.
7. Write a complete training loop from scratch, including training/validation splits, metric tracking, and device handling.
8. Build custom Datasets and DataLoaders for your own data, the way research code actually handles data.
9. Save, load, and checkpoint models so that no experiment's work is ever lost and any result can be reproduced.
10. Diagnose and fix the ten most common PyTorch errors, and organize research code with seeds, configs, and logging so your experiments are reproducible and publication-ready.

---

How to Use This Book

• Run every code block. This book is a workshop, not a lecture. Open Google Colab (Chapter 2 shows you how) and type the examples yourself. The ones you type are the ones you will remember.
• Read the "For Your Research" boxes carefully. They are the bridge between "I can run this tutorial" and "I can publish with this."
• Do the exercises. Ten are waiting at the end. Several ask you to reproduce a small experiment and log the results — exactly what reviewers will ask you to do.
• Keep the Learning Dashboard open. It summarizes every key API, shape convention, and decision rule in one place for quick reference while you code.

---

Chapter 1: Why PyTorch for Research

The framework is the lab

When a physicist starts a new project, one of the first questions is which lab instruments to use. When an AI researcher starts a new project, one of the first questions is which deep learning framework to build it in. The framework shapes everything: how fast you can test an idea, how easily others can read and reproduce your code, and whether the community's tools — pre-trained models, datasets, training utilities — plug straight into your work or fight you at every step.

There are three serious contenders today: PyTorch, TensorFlow, and JAX. All three can train neural networks. All three are used in industry. But for research — for the daily work of inventing new models, running ablations, and publishing papers — PyTorch has won, and it is worth understanding exactly why, so you can make the choice deliberately rather than by fashion.

The core difference: dynamic vs. static computation

Every deep learning framework must solve the same problem: a neural network is a long chain of math operations, and to train it you must compute gradients through that chain (backpropagation). The frameworks differ in when they build that chain.

TensorFlow (classic style) historically used a static graph: you first describe the entire computation — "take input x, multiply by matrix W, add bias b, apply ReLU" — as a graph structure, and only then do you run it with actual data. This is like writing a complete recipe before you are allowed to taste anything. It is efficient for production (the graph can be optimized, shipped to phones, served at scale), but it is awkward for research, because research is full of ideas like "what if the network's structure changes depending on the input?" or "let me print the value right here to see what's happening." In a static graph, you cannot easily print a value mid-computation or change the structure on the fly — the graph is already fixed.

PyTorch uses a dynamic graph (sometimes called "define by run"): the computation graph is built fresh every time your code runs, as a natural side effect of ordinary Python code executing. A loop in your code becomes a loop in the graph. An if statement becomes a branch in the graph. You can print any value, use a debugger, and change the architecture based on the data — because the network is just Python code running. This is the single most important reason researchers prefer PyTorch: debugging a PyTorch model feels like debugging any Python program, and experimenting feels like writing Python, not like configuring a separate graph language.

TensorFlow 2 added "eager execution" to close this gap, and JAX takes yet another approach (pure functions plus transformations like jit and grad, beloved for its mathematical elegance). But by the time the others caught up on usability, the research community had already standardized on PyTorch — and standardization itself is a powerful reason to join.

Network effects: the real reason

The honest reason most researchers use PyTorch today is not only technical. It is that everyone else does. Consider what that means in practice:

• Paper code. When authors release code with a paper (increasingly expected at top venues), it is overwhelmingly PyTorch. If you want to build on the state of the art — the single most common way to start a research project — you will be reading and modifying PyTorch code.
• Pre-trained models. The Hugging Face ecosystem, torchvision, torchaudio, and countless domain libraries distribute pre-trained models as PyTorch modules. Loading a pre-trained ResNet or BERT into your experiment is three lines of code.
• Answers to your questions. When your training loop crashes at 2 a.m., the Stack Overflow answer, the GitHub issue, and the forum thread that solve your exact problem will almost certainly be in PyTorch.
• Collaborators and reviewers. Your lab mates write PyTorch. Reviewers who ask for "code to reproduce the results" expect to run it without learning a new framework.

This is not to say TensorFlow is "bad." TensorFlow (with Keras) remains strong in production deployment, mobile, and some industry pipelines, and its SavedModel/TFLite ecosystem for shipping models is genuinely excellent. JAX is the rising star for researchers who want maximum mathematical control and speed, especially in areas like reinforcement learning and scientific machine learning — and it is worth learning later. But for your first framework, and for the fastest path into publishable research, PyTorch is the right choice. This book will note, at the end, exactly when to look at the others.

The ecosystem: what "PyTorch" actually includes

Beginners often think PyTorch is one library. In practice, "PyTorch" is a center of gravity with a large ecosystem around it:

• torch — the core: tensors, autograd, neural network modules, optimizers.
• torchvision, torchaudio, torchtext — the domain libraries: standard datasets (MNIST, CIFAR, ImageNet), image transforms, and pre-trained models. Most computer vision papers start here.
• PyTorch Lightning — a lightweight wrapper that organizes your PyTorch code into clean, reusable experiment structure (we will meet it in Chapter 12).
• Hugging Face Transformers / Diffusers / Datasets — the largest public collection of pre-trained models and datasets, built on PyTorch. If your research touches language, vision-language, or generative models, you will live here.
• Tools for the full research loop — TensorBoard and Weights & Biases for experiment tracking, torch.compile for speed, and distributed training support for multi-GPU work.

Learning core PyTorch first is the right order: every one of these tools assumes you understand tensors, autograd, modules, and the training loop. They are conveniences on top of foundations — and foundations are what this book builds.

What PyTorch asks of you

PyTorch's philosophy is often summarized as staying close to Python and NumPy. If you know NumPy, tensors will feel familiar. If you know Python classes, nn.Module will feel familiar. The framework deliberately avoids magic: there is no hidden machinery you must appease, and almost everything you need to understand is visible in your own code. That is why this book teaches PyTorch "from the inside out" — tensors first, then gradients, then modules, then the loop — instead of handing you a high-level API that hides what is happening.

One more honest note: PyTorch will let you shoot yourself in the foot. It will happily let you forget to zero your gradients, train on the wrong device, or leak memory — silently. Chapter 10 exists because every researcher, including the authors of the papers you admire, has made every one of those mistakes. Knowing the failure modes is part of knowing the framework.

For Your Research

Reading paper code is a research skill. Pick one recent paper in your area that released code (check the paper's GitHub link or Papers with Code). Before you run anything, read its training script and identify: the model definition (nn.Module), the loss function, the optimizer, and the training loop. You will find the exact four components this book teaches, in the same arrangement, in nearly every repository. Being able to map a paper's method section to these four code components is what turns "I read the paper" into "I can reproduce and extend the paper" — and that mapping is the hidden curriculum of this entire book. Start practicing it now, with Chapter 1's vocabulary: tensors, autograd, modules, losses, optimizers.

Seeing the difference: dynamic graphs in action

The dynamic-graph advantage is easiest to feel in code. Consider processing sentences of different lengths — a completely ordinary research situation. In PyTorch, you write the obvious Python:

import torch
import torch.nn as nn

rnn = nn.RNN(input_size=8, hidden_size=16, batch_first=True)

sentences = [torch.randn(1, 5, 8), torch.randn(1, 9, 8), torch.randn(1, 3, 8)]
for s in sentences:
    out, h = rnn(s)          # each forward pass builds a graph sized for THIS input
    loss = out.sum()
    loss.backward()          # gradients flow through this pass's graph
    print(s.shape[1], "->", h.shape)

No padding to a fixed length required by the framework, no graph-rebuilding ceremony, no separate "dynamic" API — the graph simply follows your code. You can put print() statements inside forward, set breakpoints, and wrap any of it in if statements driven by the data. When researchers say PyTorch is "Pythonic," this is what they mean: the distance between the idea in your head and the running code is minimal, and the distance between the running code and the debugger is zero.

When PyTorch is not the answer (honest exceptions)

A good researcher picks tools with eyes open. There are situations where the community's default shifts:

• Shipping to phones or browsers. TensorFlow Lite and ONNX Runtime have the most mature mobile/edge deployment stories. (PyTorch has TorchScript and ExecuTorch, and the gap is closing — but if your thesis requires an Android demo next month, check the deployment path first.)
• TPU training. Google's TPUs are superb for large-scale training, and JAX is their first-class citizen. PyTorch supports TPUs via XLA, but it's a second-class experience.
• Maximum-speed methods research. JAX's function transformations (jit, grad, vmap) compose in ways PyTorch is still catching up to. If your research is about differentiation itself, learn JAX as your second framework (Chapter 12).

None of these change the recommendation: learn PyTorch first, because it's where the papers, the code, and the collaborators are.

The ecosystem in one line of code

To feel the network effect concretely, this is how close "the entire field's pre-trained models" are:

from torchvision.models import resnet18, ResNet18_Weights
model = resnet18(weights=ResNet18_Weights.DEFAULT)  # pre-trained on ImageNet, ready to fine-tune

One import, one line, and you hold a model that took weeks of GPU time to train — free, documented, and in the exact nn.Module format this book teaches. That is the practical meaning of "ecosystem": other people's compute, available to your experiments in seconds.

A day in the life: the research loop PyTorch enables

To make the "why" concrete, here's what a typical research day looks like once you're fluent. Morning: you read a paper proposing a new attention variant. By lunch, you've cloned the repo, found its model.py, and located the attention layer — an nn.Module you can read like prose. Afternoon: you subclass it, tweak the forward pass, and launch an ablation with your Chapter 11 harness. Evening: the validation curves tell you whether the idea works. Total new code: maybe 60 lines.

That speed — paper to tested idea in a day — is the real product of the PyTorch ecosystem. It compounds: every paper you can read as code becomes raw material, and every experiment you can run cleanly becomes evidence. Researchers in other frameworks can do all of this too, but they do it while translating: translating paper code from PyTorch, translating answers from PyTorch forums, translating pre-trained weights into their format. Translation is a tax on every idea. PyTorch removes the tax.

Exercise for this chapter (do it now)

Pick any recent paper in your field with public code. Open its repository and find these four things: (1) the model class (nn.Module), (2) the loss function, (3) the optimizer construction, (4) the training loop. Time yourself. If you can find all four in under ten minutes, you already read PyTorch structurally — this book will just deepen that. If you can't, note where you got lost: those are exactly the chapters to focus on.

Key Takeaways

1. PyTorch's dynamic computation graph makes models ordinary Python code — easy to write, print, debug, and change mid-experiment.
2. The research community has standardized on PyTorch: paper code, pre-trained models, and community help are overwhelmingly PyTorch.
3. TensorFlow excels at production deployment; JAX excels at mathematical elegance and speed — both are worth knowing later, but PyTorch is the right first framework for research.
4. "PyTorch" means an ecosystem (torchvision, Hugging Face, Lightning, experiment trackers) built on one foundation: tensors, autograd, modules, and the training loop.
5. PyTorch gives you full control and full responsibility — understanding its failure modes is part of mastering it.

---

Chapter 2: Setup — Installation, CPU vs. GPU, and Google Colab

What you are actually installing

When you "install PyTorch," you are installing a Python package (plus its compiled math libraries) that can do fast tensor math. The single most important choice during installation is which version of the math backend you get: CPU-only, or CUDA-enabled for NVIDIA GPUs. PyTorch itself is the same Python code either way; what differs is whether operations like matrix multiplication run on your processor or on a graphics card with thousands of parallel cores.

This matters because neural network training is, at its heart, an enormous number of matrix multiplications. A GPU can do them 10 to 100 times faster than a CPU. For this book's small examples, a CPU is fine. For anything resembling real research — training on CIFAR, fine-tuning a transformer — a GPU changes the experience from "impossible" to "overnight run."

Option A: Google Colab (recommended for students)

If you are a student without a GPU machine, start with Google Colab. It is a free Jupyter-notebook service from Google that gives you access to a GPU for limited hours per day — more than enough for this book and for small research experiments. There is nothing to install: PyTorch (with GPU support) comes pre-installed.

1. Go to colab.research.google.com and sign in with a Google account.
2. Create a new notebook: File → New notebook.
3. Enable the GPU: Runtime → Change runtime type → Hardware accelerator → GPU (T4), then Save.
4. Run this in the first cell to verify:

import torch
print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("GPU name:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "none (CPU only)")

If torch.cuda.is_available() prints True, you have a working GPU. Colab's free tier gives you a modest GPU (often a T4) with usage limits — if you get a message about GPU unavailability, wait a few hours or continue on CPU; every example in this book runs on CPU too, just slower.

Colab tips that save real pain: - Files you upload or create vanish when the runtime disconnects (typically after ~12 hours, or 90 minutes idle). Save anything important to Google Drive (mount it with the file-browser button) or download it. - Install extra packages with !pip install <package> in a cell — they last for the session. - Use Runtime → Restart and run all before sharing results, so you know the notebook works top-to-bottom.

Option B: Installing on your own machine

If you have your own computer, install PyTorch with pip or conda. The official installer at pytorch.org generates the exact command for your system — use it rather than guessing, because the CUDA version must match your NVIDIA driver.

CPU-only (any laptop, no NVIDIA GPU):

pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

With NVIDIA GPU (check your CUDA version with nvidia-smi first):

# Example for CUDA 12.x — confirm the exact command at pytorch.org for your setup
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

With conda (handles CUDA libraries for you):

conda install pytorch torchvision pytorch-cuda=12.1 -c pytorch -c nvidia

After installing, verify in Python:

import torch
print(torch.__version__)          # e.g. 2.x.x
print(torch.cuda.is_available())  # True if GPU build works

A common student machine setup is: CPU-only PyTorch locally for writing and debugging code, Colab GPU for actual training runs. This is a perfectly respectable research workflow — many published experiments were trained on free Colab GPUs.

CPU vs. GPU: what actually changes in your code

Here is the good news: almost nothing. PyTorch is designed so the same code runs on both; you only choose where tensors live. The pattern you will use in every chapter from here on:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

model = MyModel().to(device)          # move the model to the device
x = torch.randn(32, 10).to(device)    # move data to the same device
out = model(x)                        # computation happens on that device

Two rules that prevent 90% of device headaches: 1. The model and the data must be on the same device. A model on GPU with data on CPU raises an error — this is the single most common beginner crash (see Chapter 10). 2. Set the device once, at the top of your script, and use it everywhere. Never hard-code "cuda".

You can also move results back to the CPU (for plotting or saving) with .cpu(), and detach them from the computation graph with .detach() — both appear constantly in training loops:

loss_value = loss.item()      # Python number, safe to log/plot
preds = outputs.detach().cpu()  # tensor on CPU, no gradient tracking

How much GPU memory do you need?

GPU memory (VRAM) is the main limit you will hit. Rough guide for this book's level:

Task	Typical VRAM needed
Toy examples in this book (small MLPs)	< 1 GB — any GPU, even CPU
CIFAR-10 with a small CNN	2–4 GB
Fine-tuning a small transformer (e.g., BERT-base)	8–12 GB
Training a transformer from scratch	16 GB+ / multiple GPUs

When you run out, you get CUDA out of memory. The standard fixes, in order: reduce batch size, use a smaller model, or use gradient accumulation (Chapter 12). Do not buy hardware for this book — Colab is enough.

A note on versions and reproducibility

PyTorch releases roughly twice a year. For research, the exact version matters less than recording it: always note the PyTorch version (and CUDA version) alongside your results. A result obtained on PyTorch 2.1 may differ in the fourth decimal place on 2.3 due to kernel changes — reviewers occasionally ask, and "I recorded it" is the entire difference between a painful week and a one-line answer. Chapter 11 makes this automatic.

For Your Research

Your environment is part of your method. Published papers increasingly include a "reproducibility" note or appendix listing software versions. Start the habit now: at the top of every experiment notebook or script, print and log torch.__version__, torch.cuda.is_available(), and the GPU name. Better yet, export your environment once per project with pip freeze > requirements.txt (or conda env export > environment.yml) and commit it with your code. When a reviewer — or your future self — asks "what exactly did you run?", the answer should be a file, not a memory.

A diagnostic script: prove your install works

After installing, don't just check the version — run a script that exercises the whole stack (tensor math, GPU, and a tiny training step). Save this as check_install.py and run it once per machine:

import torch

print("torch:", torch.__version__)
print("cuda available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("gpu:", torch.cuda.get_device_name(0))
    print("capability:", torch.cuda.get_device_capability(0))

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# tensor math on the device
x = torch.randn(1000, 1000, device=device)
print("matmul ok:", (x @ x.T).shape)

# autograd + optimizer smoke test
w = torch.randn(10, 1, device=device, requires_grad=True)
opt = torch.optim.AdamW([w], lr=1e-2)
for _ in range(5):
    opt.zero_grad()
    loss = (w ** 2).sum()
    loss.backward()
    opt.step()
print("train step ok, loss:", loss.item())
print("ALL CHECKS PASSED")

If this runs clean, your environment can do everything in this book. If it fails, the failure location tells you exactly which layer is broken (import → install problem; CUDA lines → GPU build problem; optimizer lines → deeper issue).

Common installation failures and fixes

• torch.cuda.is_available() is False on a machine with an NVIDIA GPU. Almost always a CUDA/driver mismatch: the PyTorch build's CUDA version must be ≤ your driver's supported version (nvidia-smi shows the max supported CUDA at top-right). Fix: reinstall using the exact command from pytorch.org for your CUDA version.
• pip installs the wrong variant. The default pip install torch now usually picks correctly, but on some systems it grabs CPU-only while you wanted GPU (or vice versa). Fix: use the explicit --index-url from pytorch.org, and verify with the diagnostic script.
• conda vs pip conflicts. Mixing conda-installed and pip-installed PyTorch in one environment produces bizarre import errors. Fix: pick one package manager per environment; when in doubt, create a fresh environment (conda create -n torch python=3.11 or python -m venv torch-env).
• "No module named torch" in Jupyter/Colab. The notebook kernel is using a different Python than the one you installed into. Fix: in Colab use !pip install (installs into the notebook's environment); locally, run python -c "import sys; print(sys.executable)" in both the terminal and the notebook and make them match.

Other platforms: Apple Silicon, Windows

• Apple Silicon Macs (M1/M2/M3): PyTorch supports the MPS backend — GPU acceleration via torch.device("mps"). Install the standard macOS build from pytorch.org, then use "mps" in your device selection (torch.backends.mps.is_available() to check). Most of this book works unchanged; a few exotic ops fall back to CPU automatically.
• Windows: CPU and CUDA builds work fine; the main gotcha is DataLoader(num_workers > 0) requiring the if __name__ == "__main__": guard in scripts (Chapter 8/10).

Beyond free Colab: when you need more

Free Colab is enough for this book, but real projects outgrow it. Know the ladder before you need it:

• Colab Pro/Pro+ — paid tiers with better GPUs (A100s), longer runtimes, and more memory. The cheapest way to rent serious compute by the month.
• Kaggle Notebooks — free GPU/TPU hours weekly, similar to Colab; a good second free quota when Colab's runs out.
• University clusters — if your university has one, learn its job scheduler (usually Slurm) early; it's the cheapest "big GPU" most students ever get.
• Cloud spot instances (AWS/GCP/Azure) — powerful but billed by the minute and operationally heavy; worth it only for large, well-debugged runs.

Whichever you use, the Chapter 9 + 11 habits (checkpoints to persistent storage, configs, logging) are what make cloud training survivable — a preempted spot instance without checkpoints is money burned.

Disk and environment hygiene

Two mundane things that bite everyone once:

1. Disk space. Datasets, checkpoints, and logs accumulate fast. On Colab, mount Drive and write runs/ there for anything important; locally, keep a runs/ git-ignore and clean it monthly. A run you can't find is a run you didn't do.
2. One environment per project. A single global Python with every package becomes unresolvable conflicts within months. python -m venv or conda create -n <project> per project, with its own requirements.txt — Chapter 11's layout assumes this.
3. Reinstall PyTorch deliberately. When upgrading, uninstall first (pip uninstall torch torchvision) rather than layering installs — mixed-version leftovers are a classic source of phantom import errors that waste an afternoon.

Colab survival: Drive, timeouts, and keeping your work

Colab's two harshest realities are the idle timeout (~90 minutes) and the 12-hour runtime limit — either can vaporize unsaved work. The defense is one cell at the top of every important notebook:

from google.colab import drive
drive.mount('/content/drive')

import os
WORK = "/content/drive/MyDrive/pytorch-book"
os.makedirs(f"{WORK}/checkpoints", exist_ok=True)
os.makedirs(f"{WORK}/runs", exist_ok=True)
print("work dir:", WORK)

Then point all saves at WORK: checkpoints, metrics CSVs, configs. The Chapter 9 last.pt pattern matters doubly here — set it to save every epoch into Drive, so a disconnect costs you at most one epoch. Two more Colab habits: (1) Runtime → Restart and run all before trusting any result (out-of-order cell execution is a classic source of "it worked yesterday" mysteries); (2) keep a notes.md in the Drive folder logging what each notebook run was for — Colab notebooks proliferate and their names lie.

Key Takeaways

1. Start with Google Colab if you lack a GPU: free, pre-installed PyTorch, GPU enabled in two clicks. Save important files to Drive.
2. On your own machine, use the official installer at pytorch.org to get the right CPU/CUDA build — do not guess the command.
3. CPU vs. GPU changes almost no code: pick a device once, move model and data with .to(device), and keep them on the same device.
4. GPU memory is your main constraint; the first fix for "out of memory" is a smaller batch size.
5. Record your environment (PyTorch version, CUDA, GPU) with every experiment — it is part of reproducible research.

---

Chapter 3: Tensors — The Core Data Structure

What a tensor is (and isn't)

A tensor is PyTorch's fundamental object: a multi-dimensional array of numbers. If you know NumPy, you already understand 90% of it — a PyTorch tensor is like a NumPy array with two superpowers: it can live on a GPU, and it can track operations for automatic differentiation (Chapter 4).

The word "tensor" sounds intimidating, but the idea is simple. You already know the ranks:

• A scalar (single number, e.g., a loss value) is a 0-dimensional tensor.
• A vector (a list of numbers, e.g., one data sample's features) is a 1-dimensional tensor.
• A matrix (a table of numbers, e.g., a batch of samples) is a 2-dimensional tensor.
• Anything with 3+ dimensions — an image (height × width × color channels), a batch of images, a batch of sentences — is a higher-dimensional tensor.

In deep learning, nearly everything is a tensor: your input data, the model's weights, the gradients, the outputs. Learning to think in tensors — and especially in shapes — is the single most transferable skill in this book. Most PyTorch errors are shape errors, and most shape errors come from not tracking what each dimension means.

Creating tensors

import torch

# From Python data
t1 = torch.tensor([1.0, 2.0, 3.0])                 # 1-D, shape (3,)
t2 = torch.tensor([[1, 2], [3, 4]])                # 2-D, shape (2, 2)

# Standard initializations
z = torch.zeros(2, 3)        # all zeros, shape (2, 3)
o = torch.ones(2, 3)         # all ones
r = torch.rand(2, 3)         # uniform random in [0, 1)
rn = torch.randn(2, 3)       # standard normal (mean 0, std 1) — the workhorse
eye = torch.eye(3)           # 3x3 identity matrix
seq = torch.arange(0, 10, 2) # like Python range: [0, 2, 4, 6, 8]

Two creation functions deserve special attention because they appear in almost every research script:

• torch.randn(...) — random numbers from a standard normal distribution. This is how weights are initialized (Chapter 5) and how noise is sampled in generative models.
• torch.zeros(...) / torch.ones(...) — how you pre-allocate accumulators, masks, and counters.

Every tensor has three attributes you should check constantly while learning:

x = torch.randn(32, 10)
print(x.shape)   # torch.Size([32, 10]) — the dimensions
print(x.dtype)   # torch.float32 — the number type
print(x.device)  # cpu (or cuda:0) — where it lives

dtype matters more than beginners expect. The default floating-point type is float32 — the standard for deep learning (a good balance of precision and speed/memory). Common pitfalls: creating a tensor from Python integers gives int64, and mixing int64 with float32 in an operation raises an error. When in doubt, create floats explicitly (torch.tensor([1.0, 2.0])) or convert with .float(). Model weights are float32; labels for classification are int64 (more on this in Chapter 6 — it is a deliberate convention, not an accident).

Indexing and slicing

Tensor indexing works like NumPy's, and you will use it constantly to inspect data and build masks:

x = torch.arange(12).reshape(3, 4)
# tensor([[ 0,  1,  2,  3],
#         [ 4,  5,  6,  7],
#         [ 8,  9, 10, 11]])

print(x[0])        # first row: tensor([0, 1, 2, 3])
print(x[:, 1])     # second column: tensor([1, 5, 9])
print(x[1:3, :2])  # rows 1-2, cols 0-1
print(x[x > 5])    # boolean mask: tensor([ 6,  7,  8,  9, 10, 11])

Boolean masking (x[x > 5]) is worth practicing: it is how you filter predictions, select hard examples, and implement many research ideas (like focusing loss on misclassified samples) in one line.

Shape manipulation: the daily bread

Reshaping tensors is something you will do in nearly every model. The key operations:

x = torch.arange(12)            # shape (12,)
m = x.reshape(3, 4)             # shape (3, 4) — same data, new shape
m2 = x.view(3, 4)               # like reshape, but requires contiguous memory
t = m.T                         # transpose: shape (4, 3)
flat = m.reshape(-1)            # -1 means "infer this dimension": shape (12,)
batched = x.reshape(2, 2, 3)    # shape (2, 2, 3)

reshape vs view: both change shape without copying data when possible. view fails if the tensor's memory layout doesn't allow it; reshape falls back to copying. Rule of thumb: use reshape unless you have a reason to demand no-copy (view), and you will rarely go wrong.

Three more shape tools that appear everywhere in real code:

x = torch.randn(32, 10)          # a batch of 32 samples, 10 features each
x2 = x.unsqueeze(0)              # add a dim at front: (1, 32, 10)
x3 = x.unsqueeze(-1)             # add a dim at end: (32, 10, 1)
back = x3.squeeze(-1)            # remove size-1 dims: (32, 10)
cat = torch.cat([x, x], dim=0)   # concatenate: (64, 10)
stack = torch.stack([x, x])      # stack as new dim: (2, 32, 10)

Learn the difference between cat and stack cold: cat glues along an existing dimension (two batches become one bigger batch), stack creates a new dimension (a list of batches becomes one batched tensor). Confusing them is a classic shape bug.

Broadcasting: the silent helper (and silent footgun)

Broadcasting is PyTorch's rule for operating on tensors of different shapes: dimensions are aligned from the right, and any dimension of size 1 (or a missing dimension) is stretched to match. Examples:

a = torch.ones(3, 4)
b = torch.ones(4)        # shape (4,) aligns with last dim -> result (3, 4)
c = a + b                # each row of a gets b added

m = torch.ones(3, 1)
n = torch.ones(1, 4)
print((m + n).shape)     # (3, 4) — both stretched

Broadcasting is why you can write x - mean to normalize a whole batch by one mean vector, or add a bias vector to every row of a batch. It is elegant — and it is also the source of bugs where shapes almost match and PyTorch silently does something you didn't intend. The defense is a habit: whenever a result looks wrong, print shapes first. A one-line print(x.shape) has solved more deep learning bugs than any debugger.

Essential operations

Element-wise ops work as you'd expect (+, -, *, /, **, comparisons). The two matrix operations to never confuse:

a = torch.randn(3, 4)
b = torch.randn(4, 5)

elem = a * a              # element-wise: (3, 4)
mat = a @ b               # matrix multiplication: (3, 5)
mat2 = torch.matmul(a, b) # same as @

@ is matrix multiplication; * is element-wise. A neural network layer is essentially @ plus a bias — so this distinction is deep learning arithmetic.

Reductions collapse dimensions and are how losses and metrics are computed:

x = torch.randn(32, 10)
print(x.sum())            # single number: sum of everything
print(x.mean(dim=0).shape) # (10,) — mean over the batch: per-feature means
print(x.max(dim=1))       # values and indices of max per row
print(x.argmax(dim=1).shape) # (32,) — predicted class per sample (classification!)

argmax(dim=1) deserves a mental bookmark: "for each sample in the batch, which class scored highest" is exactly how you turn model outputs into predictions (Chapter 7).

Moving between devices, and the NumPy bridge

x = torch.randn(100, 50)
if torch.cuda.is_available():
    x = x.to("cuda")        # or x.cuda()
    y = x @ x.T             # runs on GPU
    x = x.to("cpu")         # or x.cpu()

# NumPy interop (CPU tensors only) — shares memory, no copy!
import numpy as np
arr = np.array([1.0, 2.0, 3.0])
t = torch.from_numpy(arr)   # tensor shares memory with arr
back = t.numpy()            # back to NumPy, also shared

The memory-sharing is a double-edged sword: modifying one modifies the other. It is efficient (no copies), but be deliberate about it. Also note: .numpy() only works on CPU tensors — move to CPU first.

Worked example: normalizing a batch, the tensor way

Here is a realistic mini-task — standardizing features (zero mean, unit variance), written purely with tensor ops:

torch.manual_seed(0)
data = torch.randn(1000, 5) * 10 + 3   # 1000 samples, 5 features, arbitrary scale

mean = data.mean(dim=0)                # (5,) — per-feature mean
std = data.std(dim=0)                  # (5,) — per-feature std
normalized = (data - mean) / std       # broadcasting over the batch dim

print(normalized.mean(dim=0))  # ≈ [0, 0, 0, 0, 0]
print(normalized.std(dim=0))   # ≈ [1, 1, 1, 1, 1]

Three lines, no loops. This "think in whole tensors, not loops" style is the core performance habit of PyTorch: loops over samples in Python are slow; tensor ops run in optimized C++/CUDA. Whenever you catch yourself writing for i in range(batch_size), stop and look for the tensor op — it almost always exists.

For Your Research

Shape discipline is paper discipline. Every tensor in your experiment has a meaning per dimension: (batch, channels, height, width) for images, (batch, sequence, features) for sequences. Write these conventions in a comment at the top of your data pipeline and model — e.g., # x: (B, C, H, W). Reviewers and collaborators read shape conventions like notation in a math paper; inconsistent or undocumented shapes are the code equivalent of undefined symbols. When you later write your paper's method section, these shape annotations become your equations' dimensions almost verbatim.

Memory layout: why .view() sometimes fails

Tensors have a hidden property: memory layout — the order elements are stored in. A freshly created (3, 4) tensor stores its 12 numbers row by row (row-major, "contiguous"). reshape and view only avoid copying when the new shape is compatible with that layout. Transposing scrambles the layout:

x = torch.arange(12).reshape(3, 4)
print(x.is_contiguous())        # True
t = x.T
print(t.is_contiguous())        # False — columns are now strided
# t.view(12)                    # RuntimeError: view size is not compatible
print(t.reshape(12).shape)      # (12,) — reshape copies silently and works
print(t.contiguous().view(12).shape)  # (12,) — explicit copy, then view works

Practical rule: use reshape by default (it does the right thing either way). Reach for view only when you specifically need the no-copy guarantee, and call .contiguous() first when a view fails after a transpose or permute. This is one of those errors that looks terrifying ("view size is not compatible with input tensor's size and stride") but has a one-word fix once you understand layout.

In-place operations: speed with strings attached

Appending _ to an op makes it in-place (relu_(), add_(), mul_()), modifying the tensor instead of allocating a new one. This saves memory — meaningful for giant models — but has two costs: (1) it breaks autograd if the tensor is needed for gradient computation (Chapter 10, Error 8), and (2) it makes debugging harder because values are destroyed. Guidance: never use in-place ops while developing; consider them only as a late optimization for memory-bound models, and never on tensors that require gradients.

Advanced indexing: gather, masked fill, where

Three more indexing tools that appear constantly in research code:

x = torch.randn(4, 5)

# torch.where: element-wise if-else — enormously useful
y = torch.where(x > 0, x, torch.zeros_like(x))  # = ReLU, written manually

# masked_fill: set masked positions to a value (attention mechanisms live on this)
mask = torch.tensor([True, False, True, False, True])
row = torch.randn(5)
print(row.masked_fill(mask, float("-inf")))  # masked positions become -inf

# gather: pick one element per row according to indices (used in many losses)
logits = torch.randn(4, 10)
labels = torch.tensor([3, 0, 9, 1])
picked = logits.gather(1, labels.unsqueeze(1))  # shape (4, 1): the logit of the true class

gather deserves a second look: "for each sample, get the score of its true class" is the inner operation of cross-entropy loss itself. When you read advanced loss implementations, you'll find gather doing the precise selection work.

A note on precision: float32 vs float16

Deep learning overwhelmingly uses float32 (the default). float64 doubles memory for no accuracy benefit in almost all cases — don't use it. float16/bfloat16 (half precision) halve memory and speed up modern GPUs, but need care (loss scaling) — that's Chapter 12's mixed-precision topic. For this book: stay in float32 unless you have a measured reason to leave.

Key Takeaways

1. A tensor is a multi-dimensional array with GPU support and gradient tracking — the universal currency of PyTorch.
2. Always know a tensor's shape, dtype, and device; print them whenever something looks wrong.
3. reshape/view, unsqueeze/squeeze, cat/stack are your daily shape tools — learn cat vs stack precisely.
4. Broadcasting aligns shapes from the right and stretches size-1 dims: powerful for batch ops, dangerous when shapes almost match.
5. @ is matrix multiplication, * is element-wise — never confuse them; think in whole-tensor ops, not Python loops.
6. Tensors bridge to NumPy with shared memory (CPU only), and move between devices with .to(device).

---

Chapter 4: Autograd — Automatic Differentiation Explained Gently

The problem autograd solves

Training a neural network means adjusting millions of numbers (weights) to make the loss go down. To know how to adjust each weight, you need the gradient: how much would the loss change if this weight changed a tiny bit? Computing gradients for millions of weights by hand is impossible; computing them by finite differences (nudge each weight, re-run, measure) would take millions of forward passes.

Autograd — automatic differentiation — is PyTorch's solution. You write the forward computation normally, and PyTorch silently records every operation in a computation graph. When you call .backward(), it walks that graph in reverse, applying the chain rule from calculus, and fills in every gradient in a single pass. One forward pass, one backward pass, all gradients. This is backpropagation, automated.

The beautiful part: you never write derivative code. You only write the forward math — the thing you understand — and the framework handles the calculus.

The minimal example

import torch

x = torch.tensor(2.0, requires_grad=True)  # "track operations on x"
y = x ** 2 + 3 * x + 1                     # y = x² + 3x + 1; graph recorded
print(y)  # tensor(11., grad_fn=<AddBackward0>)

y.backward()      # compute dy/dx via chain rule
print(x.grad)     # tensor(7.) — because dy/dx = 2x + 3 = 7 at x=2

Three things happened: 1. requires_grad=True told PyTorch to record operations involving x. 2. Computing y built a graph of operations (pow, mul, add) behind the scenes — note grad_fn in the printout, which is the graph node. 3. y.backward() traversed the graph backward, computing the gradient of y with respect to x and storing it in x.grad.

Check the math: y = x² + 3x + 1, so dy/dx = 2x + 3. At x = 2: 4 + 3 = 7. ✓ The framework did calculus correctly, and it will do it correctly for a graph with a million nodes too.

Gradients of vectors: backward needs a scalar

.backward() computes gradients of a scalar output. If your output is a vector, you must reduce it to a scalar first (usually with .sum()), or pass explicit gradient weights. In practice this is never a problem: loss functions always produce a scalar (one number summarizing "how wrong"), so you always call loss.backward().

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2
# y.backward()          # ERROR: y is not a scalar
y.sum().backward()      # fine: gradient of sum(x²)
print(x.grad)           # tensor([2., 4., 6.]) — d(x²)/dx = 2x

The computation graph is rebuilt every time (dynamic!)

This is where Chapter 1's "dynamic graph" pays off. The graph is created fresh on each forward pass, so it can depend on data, loops, and conditions:

x = torch.tensor(2.0, requires_grad=True)
y = x
for _ in range(3):
    y = y * x + 1      # loop unrolled into the graph, 3 iterations
y.backward()
print(x.grad)  # correct gradient through the loop — try verifying by hand!

Try doing that with a static graph framework's old API — you would be fighting the framework. In PyTorch it is just Python. This dynamic nature is exactly what lets researchers invent architectures with data-dependent control flow (like attention mechanisms that loop over variable-length sequences).

You can inspect the graph if you're curious — each tensor with a grad_fn points to the operation that created it, which points to its inputs:

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
z = y + 1
print(z.grad_fn)            # <AddBackward0>
print(z.grad_fn.next_functions)  # shows the PowBackward node feeding it

You will rarely inspect graphs manually, but knowing they exist — and are rebuilt per forward pass — explains several behaviors below.

Gradient accumulation: the #1 beginner trap

Gradients accumulate. Calling .backward() adds to .grad; it does not overwrite. This is deliberate (it enables training on large effective batches by accumulating over mini-batches), but it means in a training loop you must zero gradients each step:

x = torch.tensor(2.0, requires_grad=True)
for step in range(3):
    x.grad = None            # or optimizer.zero_grad() in real code
    y = x ** 2
    y.backward()
    print(f"step {step}: grad = {x.grad.item()}")  # 4.0 every time

Without the zeroing line, you'd see 4.0, 8.0, 12.0 — the classic "my loss explodes / behaves weirdly" bug. Burn this into memory: zero grads, forward, backward, step — the four-beat rhythm of every training loop (Chapter 7).

(In modern code you will call optimizer.zero_grad(), which does exactly this for all model parameters.)

Controlling tracking: no_grad and detach

Not every computation needs gradients. During evaluation (measuring accuracy on test data), tracking gradients wastes memory and time. Two tools:

x = torch.tensor(2.0, requires_grad=True)

with torch.no_grad():          # nothing inside is tracked
    y = x ** 2
print(y.requires_grad)         # False

z = (x ** 2).detach()          # detach: cut this tensor out of the graph
print(z.requires_grad)         # False

• torch.no_grad() — a context manager: "don't build graphs in here." Use it around all evaluation/inference code.
• .detach() — returns a tensor sharing the same data but cut off from the graph. Use it when you need the values (for logging, plotting, metrics) without the gradients.

A subtle one: torch.no_grad() is also how you do manual weight updates safely (otherwise the update itself gets tracked and the graph grows forever — a real memory-leak bug).

requires_grad on parameters, and leaf nodes

In a neural network, you set requires_grad=True on the weights (the things you want to learn), not on the input data:

w = torch.randn(3, requires_grad=True)   # weights: learn these
b = torch.tensor(0.0, requires_grad=True)
x = torch.tensor([1.0, 2.0, 3.0])        # data: no grad needed

pred = (w * x).sum() + b
loss = (pred - 10.0) ** 2
loss.backward()
print(w.grad, b.grad)  # gradients for w and b; x.grad is None

Note x.grad is None — PyTorch only accumulates gradients into leaf tensors with requires_grad=True (tensors you created directly, not results of operations). Intermediate values' gradients are computed during the backward pass but discarded unless you ask to keep them (tensor.retain_grad() — occasionally useful when debugging).

A complete worked example: fitting a line by hand

Let's tie it together — linear regression with explicit autograd, no nn.Module yet (that's Chapter 5). This is the smallest possible "training":

import torch

torch.manual_seed(0)
# True relationship: y = 3x + 2, plus noise
x = torch.rand(100, 1) * 10
y = 3 * x + 2 + torch.randn(100, 1)

w = torch.randn(1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
lr = 0.01

for epoch in range(200):
    w.grad = None; b.grad = None     # 1. zero grads
    pred = w * x + b                 # 2. forward
    loss = ((pred - y) ** 2).mean()  # 3. scalar loss (MSE)
    loss.backward()                  # 4. backward: gradients into w.grad, b.grad
    with torch.no_grad():            # 5. update weights (no tracking!)
        w -= lr * w.grad
        b -= lr * b.grad

print(f"learned: w={w.item():.3f}, b={b.item():.3f}")  # ≈ w=3.0, b=2.0
print(f"final loss: {loss.item():.4f}")

Run it: w converges to ~3 and b to ~2. You just did machine learning with five lines of loop. Every training loop in this book — and in published research — is this same pattern with bigger models and better tooling: zero grads → forward → loss → backward → step. When training code confuses you later, come back to this example and find the five steps.

What autograd cannot do (honest limits)

• It differentiates through the operations you ran. Non-differentiable ops (like argmax, or converting to a Python int mid-graph) break the chain — gradients won't flow through them. (This is why classification uses softmax + cross-entropy instead of differentiating through "pick the max.")
• In-place operations on tracked tensors (like x += 1) can corrupt the graph — PyTorch will raise an error rather than give wrong gradients. Prefer out-of-place ops (x = x + 1).
• It needs the graph in memory during backward. Very deep models on long sequences can run out of memory — the standard fix is gradient checkpointing (Chapter 12).

For Your Research

Autograd is why you can invent new layers. Any differentiable computation you can write in PyTorch becomes trainable automatically. When you read a paper proposing a novel layer or loss, the authors did not derive gradients by hand — they wrote the forward computation and let autograd handle the rest. This has a practical consequence for your own work: when you have an idea for a custom loss or architecture tweak, prototype the forward version first and check that loss.backward() runs and gradients are non-zero/non-NaN. If the forward is correct and differentiable, the backward is free. Many "novel" paper contributions are exactly this: a clever forward computation plus autograd.

The chain rule with real numbers

The chain rule is the entire mathematical content of backpropagation, and it's worth seeing once with concrete numbers. Take y = (2x + 1)² at x = 3. Break it into steps: a = 2x + 1 = 7, then y = a² = 49. The chain rule says: dy/dx = (dy/da) × (da/dx) = 2a × 2 = 14 × 2 = 28.

Autograd does exactly this, mechanically, for millions of chained steps:

x = torch.tensor(3.0, requires_grad=True)
a = 2 * x + 1     # a = 7,   da/dx = 2
y = a ** 2        # y = 49,  dy/da = 2a = 14
y.backward()
print(x.grad)     # 28.0 = 14 * 2 ✓

Every .backward() call is this same multiplication of local derivatives, chained from the loss back to every weight. There is no magic — just the chain rule applied at a scale no human could do by hand. When someone says "backprop," they mean this walk; when they say "autograd does backprop," they mean PyTorch builds the chain for you.

Retaining gradients and higher-order derivatives

By default, only leaf tensors (like weights) keep their .grad; intermediate gradients are freed during the backward pass to save memory. If you're debugging and want to see an intermediate's gradient, call retain_grad() before backward:

x = torch.tensor(2.0, requires_grad=True)
h = (x ** 2).retain_grad() if False else x ** 2  # placeholder — see below

More precisely:

x = torch.tensor(2.0, requires_grad=True)
h = x ** 2
h.retain_grad()          # ask to keep h's gradient
y = h ** 3
y.backward()
print(h.grad)            # d(h³)/dh = 3h² = 3*16 = 48

PyTorch can also differentiate through a backward pass (create_graph=True), giving second derivatives — used in meta-learning and some optimization research. You won't need it in this book, but knowing it exists demystifies papers that mention "higher-order gradients": it's the same machinery, applied twice.

A peek under the hood: custom autograd functions

Rarely, you need an operation autograd doesn't know how to differentiate (a novel activation, a custom CUDA kernel). PyTorch lets you define both directions explicitly:

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)       # stash what backward needs
        return x.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        return grad_out * (x > 0).float()  # gradient flows only where x > 0

x = torch.tensor([-1.0, 2.0], requires_grad=True)
y = MyReLU.apply(x).sum()
y.backward()
print(x.grad)  # tensor([0., 1.]) ✓ — correct ReLU gradients

You will likely never write one of these in your first year — but when you read a paper whose repository contains a functions.py with autograd.Function subclasses, you'll now recognize exactly what it's doing: hand-specifying the forward and its derivative for an operation the authors invented.

Key Takeaways

1. Autograd records operations into a computation graph during the forward pass; .backward() applies the chain rule to fill in all gradients at once.
2. requires_grad=True marks what to track (usually weights); .backward() requires a scalar output (your loss).
3. The graph is rebuilt every forward pass — loops and conditionals just work.
4. Gradients accumulate — zero them every step (optimizer.zero_grad()), or training silently breaks.
5. Use torch.no_grad() for evaluation and .detach() when you need values without gradients.
6. The training rhythm is always: zero grads → forward → loss → backward → step.

---

Chapter 5: Your First Neural Network with nn.Module

From raw tensors to modules

In Chapter 4 you trained a linear model with hand-managed tensors w and b. That works for two parameters. It does not work for two million. Real networks need three things hand-rolled tensors don't provide: organized parameters (so an optimizer can find and update all of them), composable layers (so you can stack building blocks), and behavior modes (training vs. evaluation, which affects layers like dropout).

torch.nn.Module is the base class that provides all three. Everything in PyTorch — a single linear layer, a ResNet, a transformer — is an nn.Module. Learning this one class deeply is learning the framework's entire modeling system.

The anatomy of a module

An nn.Module subclass has two parts you write, and machinery you inherit:

import torch
import torch.nn as nn

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()                 # always call this first
        self.fc1 = nn.Linear(10, 32)       # layer: 10 inputs -> 32 outputs
        self.relu = nn.ReLU()              # activation function
        self.fc2 = nn.Linear(32, 1)        # layer: 32 inputs -> 1 output

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = TinyNet()
print(model)

__init__ declares the sub-modules (layers); forward describes how data flows through them. You never call model.forward(x) directly — you call model(x). The __call__ machinery does important bookkeeping around your forward (hooks, mode flags), so always call the module, never forward.

nn.Linear(in_features, out_features) is the workhorse layer: it holds a weight matrix of shape (out_features, in_features) and a bias vector of shape (out_features), and computes x @ W.T + b. Its weights are initialized sensibly by default (a uniform distribution scaled by layer size — a variant of the Xavier/Glorot initialization from reference [9]).

Parameters: the things that learn

Every tensor inside a module that should be learned is a parameter. Modules track them automatically:

model = TinyNet()
for name, p in model.named_parameters():
    print(name, p.shape, p.requires_grad)
# fc1.weight torch.Size([32, 10]) True
# fc1.bias   torch.Size([32]) True
# fc2.weight torch.Size([1, 32]) True
# fc2.bias   torch.Size([1]) True

n_params = sum(p.numel() for p in model.parameters())
print("total parameters:", n_params)  # 32*10+32 + 1*32+1 = 385

model.parameters() is what you hand to the optimizer (Chapter 6): "these are the numbers to adjust." Counting parameters with the one-liner above is a habit worth forming — papers report parameter counts, and yours should match your claims.

If you need a custom learnable tensor (not inside a standard layer), wrap it in nn.Parameter:

class ScaleShift(nn.Module):
    def __init__(self):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(1))   # learnable!
        self.shift = nn.Parameter(torch.zeros(1))  # learnable!
    def forward(self, x):
        return x * self.scale + self.shift

A plain torch.tensor assigned as an attribute would not be registered — the optimizer would silently ignore it. This is a real bug researchers hit: "my custom parameter isn't learning." The fix is always nn.Parameter.

Activation functions: why networks need them

Without activations, stacking linear layers is pointless — a stack of linear operations collapses into one linear operation, and the network could only ever learn straight-line relationships. Activation functions apply a non-linearity after each layer, which is what lets networks learn curves, boundaries, and complex patterns.

The ones you need first:

relu = nn.ReLU()        # max(0, x) — the default hidden-layer activation
sigmoid = nn.Sigmoid()  # 1/(1+e^-x) — squashes to (0, 1); binary outputs
tanh = nn.Tanh()        # squashes to (-1, 1); older hidden activation
softmax = nn.Softmax(dim=1)  # turns scores into probabilities summing to 1

ReLU (max(0, x)) is the standard for hidden layers: simple, fast, and it avoids the vanishing-gradient problems of sigmoid/tanh in deep stacks. You will use it 95% of the time in this book. Softmax appears at classification outputs (with cross-entropy loss, Chapter 6). Sigmoid appears for binary decisions and gating.

nn.Sequential: the shortcut (and its limits)

For a plain stack of layers, nn.Sequential is cleaner:

model = nn.Sequential(
    nn.Linear(10, 32),
    nn.ReLU(),
    nn.Linear(32, 16),
    nn.ReLU(),
    nn.Linear(16, 1),
)
out = model(torch.randn(8, 10))  # batch of 8 -> (8, 1)
print(out.shape)

Use Sequential for simple stacks; write a custom Module class when you need branches, skip connections, multiple inputs/outputs, or any logic in forward. Research models are almost always custom classes — ResNet's skip connections can't be expressed as a plain sequence.

Worked example: a classifier on synthetic data

Let's build a complete small classifier — two-moons-style data, a 2-layer network, trained with the Chapter 4 loop pattern. (Chapters 6–7 will give this proper tooling; here the goal is to see nn.Module in action.)

import torch
import torch.nn as nn

torch.manual_seed(42)

# --- synthetic 2-D binary classification data: two blobs ---
n = 200
blob0 = torch.randn(n, 2) * 0.8 + torch.tensor([-2.0, 0.0])
blob1 = torch.randn(n, 2) * 0.8 + torch.tensor([2.0, 0.0])
X = torch.cat([blob0, blob1])                       # (400, 2)
y = torch.cat([torch.zeros(n, 1), torch.ones(n, 1)])  # (400, 1)

class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 16), nn.ReLU(),
            nn.Linear(16, 16), nn.ReLU(),
            nn.Linear(16, 1),            # single logit output
        )
    def forward(self, x):
        return self.net(x)

model = Classifier()
opt = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.BCEWithLogitsLoss()   # binary classification loss (Chapter 6)

for epoch in range(300):
    opt.zero_grad()
    logits = model(X)                # (400, 1)
    loss = loss_fn(logits, y)
    loss.backward()
    opt.step()
    if epoch % 100 == 0:
        print(f"epoch {epoch}: loss={loss.item():.4f}")

with torch.no_grad():
    preds = (torch.sigmoid(model(X)) > 0.5).float()
    acc = (preds == y).float().mean()
print(f"train accuracy: {acc.item():.3f}")  # expect ~1.000 on this easy data

Notice how the pieces slot together: nn.Module holds the architecture and parameters, model.parameters() feeds the optimizer, model(X) runs the forward pass, and the Chapter 4 rhythm (zero → forward → loss → backward → step) drives learning. You now have the complete vocabulary of a PyTorch training script.

Train vs. eval mode: a preview that prevents real bugs

Modules have two modes, switched with model.train() and model.eval(). They matter for layers whose behavior differs between training and inference — dropout (randomly zeroes activations during training only) and batch normalization (uses batch statistics in training, running statistics in eval). Our tiny nets above don't use these layers, so the modes change nothing yet — but the habit must be built now:

model.train()   # before the training loop
# ... training ...
model.eval()    # before evaluation — and combine with torch.no_grad()

Forgetting model.eval() before measuring test accuracy is a classic source of "my test results are noisy / don't reproduce." Chapter 7 makes it automatic; Chapter 10 explains the failure.

For Your Research

Your model class is the executable version of your paper's method section. When reviewers read "we use a 3-layer MLP with 128 hidden units and ReLU activations," they will check your nn.Module and expect to find exactly that. Write model classes that mirror your paper's description one-to-one: name layers after the paper's notation where possible, and put the architectural hyperparameters (hidden sizes, depth, dropout rates) in __init__ arguments rather than hard-coding them. A model class whose constructor reads MyModel(hidden_dim=128, num_layers=3, dropout=0.1) is self-documenting, easy to ablate (Chapter 11), and trivially maps to a methods paragraph.

Initialization: why starting values matter

nn.Linear doesn't start with random junk — it uses a scaled uniform distribution (Kaiming/He-style for the default). This matters enormously: initialize too large and activations explode through deep stacks; too small and signals vanish. The defaults are good, but you should know how to take control:

import torch.nn as nn

layer = nn.Linear(64, 32)
nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu")  # matched to ReLU
nn.init.zeros_(layer.bias)

# Xavier/Glorot (from reference [9]) — matched to tanh/sigmoid:
# nn.init.xavier_uniform_(layer.weight)

The principle (Glorot & Bengio [9]; He et al. for ReLU): keep the variance of activations stable as signals pass through layers. If you ever build a custom layer or see "my deep network won't train at all," initialization is on the suspect list — and the one-line fix is using the initializer matched to your activation.

Dropout and BatchNorm: layers with two personalities

Two layers that make model.train() / model.eval() actually matter:

model = nn.Sequential(
    nn.Linear(64, 64),
    nn.BatchNorm1d(64),   # normalizes activations using batch statistics
    nn.ReLU(),
    nn.Dropout(0.3),      # randomly zeroes 30% of activations (train only)
    nn.Linear(64, 10),
)

model.train()  # Dropout active; BatchNorm uses current batch stats
model.eval()   # Dropout off; BatchNorm uses learned running stats

• Dropout (Srivastava et al. [8]) randomly zeroes a fraction of activations during training, forcing the network not to rely on any single neuron — a cheap, effective regularizer. At eval it's a no-op.
• BatchNorm (Ioffe & Szegedy [7]) normalizes each layer's inputs using batch statistics during training (and tracks running averages), switching to those running averages at eval. It stabilizes and speeds up training of deep networks dramatically.

The classic beginner bug is evaluating with model.train() still on: dropout randomness makes predictions noisy and BatchNorm uses single-batch stats — your "test accuracy" becomes a random number. The eval ritual from Chapter 7 exists for exactly these layers.

Freezing layers: the transfer-learning pattern

When fine-tuning a pre-trained model (Chapter 9/12), you often freeze early layers (keep their learned features) and train only the head:

from torchvision.models import resnet18, ResNet18_Weights
model = resnet18(weights=ResNet18_Weights.DEFAULT)

for p in model.parameters():
    p.requires_grad = False          # freeze everything...
model.fc = nn.Linear(model.fc.in_features, 10)  # ...except a new head
# (the new layer's parameters default to requires_grad=True)

opt = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3)

requires_grad = False removes those parameters from the graph — no gradients computed, no updates, less memory. The optimizer only sees the unfrozen head. This pattern — frozen backbone, trained head — is behind a huge fraction of applied deep learning papers, and it's five lines once you understand parameters.

Buffers vs. parameters: state that doesn't learn

Not everything a module holds should be learned. BatchNorm's running mean/variance, for example, are updated during training but not by gradient descent — they're buffers, registered with register_buffer:

class MyNorm(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(n))   # learned by optimizer
        self.register_buffer("running_mean", torch.zeros(n))  # updated manually, saved in state_dict

    def forward(self, x):
        if self.training:
            batch_mean = x.mean(dim=0)
            self.running_mean.lerp_(batch_mean, 0.1)  # moving average update
            return (x - batch_mean) * self.weight
        return (x - self.running_mean) * self.weight

Key properties of buffers: they appear in state_dict (so they save/load with the model), they move with .to(device), but the optimizer ignores them. When you read library code and see register_buffer, that's the signal: persistent state, not learned weights. A common example you'll meet: positional encodings in transformers are often buffers.

model.apply: surgery on every submodule

Need to initialize (or modify) every layer of a deep network at once? model.apply(fn) recursively visits every submodule:

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_uniform_(m.weight, nonlinearity="relu")
        nn.init.zeros_(m.bias)

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 5))
model.apply(init_weights)   # visits every Linear, including nested ones

This is how research code applies custom initialization, freezes/unfreezes layer types, or converts layers across a whole model in one shot. Combined with named_parameters() filtering (Chapter 5's freezing pattern), it gives you complete programmatic control over large architectures.

ModuleList / ModuleDict: the list trap

A nasty gotcha: submodules in a plain Python list are invisible to the module system — their parameters won't appear in model.parameters(), won't move with .to(device), and won't save in the state_dict:

class Broken(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(10, 10) for _ in range(3)]  # BUG: plain list!

class Fixed(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(3)])  # registered

print(sum(p.numel() for p in Broken().parameters()))  # 0 — nothing to train!
print(sum(p.numel() for p in Fixed().parameters()))   # 330 — all registered

nn.ModuleList (ordered) and nn.ModuleDict (named) exist for exactly this: use them whenever layers are built in a loop or stored in a container. The symptom of the bug — "loss never decreases, gradients are None" — is maddening until you've seen it once; now you have.

Key Takeaways

1. nn.Module is the universal base class: declare layers in __init__, data flow in forward, and always call model(x), never model.forward(x).
2. model.parameters() exposes everything learnable — hand it to the optimizer; count parameters with sum(p.numel() for p in model.parameters()).
3. Custom learnable tensors must be wrapped in nn.Parameter, or the optimizer will silently ignore them.
4. ReLU is the default hidden activation; softmax/sigmoid belong at outputs for classification.
5. nn.Sequential for plain stacks; custom classes for anything with branches or skip connections.
6. Build the model.train() / model.eval() habit now — it prevents real evaluation bugs later.

---

Chapter 6: Loss Functions and Optimizers — How to Choose

Two decisions that shape every experiment

Every training run is driven by two choices you make before the loop starts:

1. The loss function — a single number measuring "how wrong is the model right now?" It defines what it means to be right, so it defines what the model learns.
2. The optimizer — the algorithm that converts gradients into weight updates. It defines how the model moves toward being right.

Choose the loss for the task; choose the optimizer for the training dynamics. This chapter gives you the standard pairings, the reasoning behind them, and the decision tables for your dashboard.

Loss functions: match the loss to the task

Regression (predicting a number): Mean Squared Error

loss_fn = nn.MSELoss()
loss = loss_fn(predictions, targets)  # mean((pred - target)^2)

MSE penalizes large errors quadratically — being off by 4 is sixteen times worse than being off by 1. That makes it sensitive to outliers; if your targets have extreme values, consider nn.L1Loss (mean absolute error, robust to outliers) or nn.SmoothL1Loss (a blend: quadratic near zero, linear far away — the default in many detection papers).

Binary classification (yes/no): BCEWithLogitsLoss

loss_fn = nn.BCEWithLogitsLoss()   # input: raw logits (N, 1); target: 0/1 floats (N, 1)
loss = loss_fn(model(X), y.float())

The name says it all: binary cross-entropy, computed from logits (raw scores, no sigmoid applied). It applies the sigmoid internally in a numerically stable way. Do not apply sigmoid yourself before this loss — that double-squashes and silently degrades training. (Chapter 10's error list includes this one.)

Multi-class classification (which of K classes): CrossEntropyLoss

loss_fn = nn.CrossEntropyLoss()    # input: logits (N, K); target: class indices (N,) as int64
loss = loss_fn(model(X), y_long)   # y contains 0..K-1, NOT one-hot, NOT probabilities

This is the most-used loss in deep learning, and the most misunderstood by beginners. Three facts to memorize: 1. Input is raw logits — no softmax beforehand (like BCE, softmax is applied internally for stability). 2. Targets are class indices (torch.long), shape (N,) — not one-hot vectors, not floats. 3. It combines LogSoftmax + negative log-likelihood: it rewards the model for assigning high probability to the correct class.

If your targets are one-hot encoded, either convert with y.argmax(dim=1) or use nn.NLLLoss with a manual LogSoftmax — but the standard pipeline is indices + CrossEntropyLoss.

The decision table:

Task	Output shape	Target format	Loss
Regression	(N, 1) or (N,)	floats	nn.MSELoss (or L1Loss/SmoothL1Loss)
Binary classification	(N, 1) logits	float 0/1, (N, 1)	nn.BCEWithLogitsLoss
Multi-class (K classes)	(N, K) logits	int64 indices (N,)	nn.CrossEntropyLoss
Multi-label (several tags)	(N, K) logits	float 0/1 (N, K)	nn.BCEWithLogitsLoss

Optimizers: how weights actually move

The optimizer implements the update rule. You create it with the model parameters and a learning rate, then call zero_grad() / step() around the backward pass:

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

optimizer.zero_grad()   # clear old gradients
loss.backward()         # compute new gradients
optimizer.step()        # update weights: w -= lr * (something like gradient)

SGD (stochastic gradient descent) — the original, from Robbins & Monro [11]:

opt = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)

Plain SGD steps directly downhill. Momentum (typically 0.9) adds a "velocity" that smooths out noisy mini-batch gradients and speeds through ravines — almost always enabled. Weight decay adds a gentle pull toward small weights (L2 regularization), fighting overfitting. SGD + momentum remains the choice for many vision papers (ResNets were trained this way [6]) because it often generalizes slightly better, though it needs more learning-rate tuning.

Adam — the default starting point (Kingma & Ba [2]):

opt = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)

Adam keeps a per-parameter learning rate: it adapts using running averages of the gradient (momentum-like) and of the squared gradient (scaling down updates for frequently-large gradients). In practice: Adam converges faster and is far less sensitive to the initial learning rate, which is why it's the default for most new projects, transformers, and quick experiments. Its downsides: slightly worse generalization than well-tuned SGD in some vision tasks, and the original Adam interacts oddly with L2 regularization — which is why AdamW (Loshchilov & Hutter [10]) exists and is now the standard Adam variant: it decouples weight decay from the adaptive update. Rule of thumb: use AdamW instead of Adam for new work.

The decision table:

Situation	Optimizer	Typical LR	Notes
Default / new project / transformer	AdamW	1e-3 (1e-4–5e-4 for transformers)	Robust, fast convergence
CNN from scratch (vision papers)	SGD + momentum 0.9	0.1 (with schedule)	Often best final accuracy, needs tuning
Quick prototype, small MLP	Adam/AdamW	1e-3	Least fiddling
Fine-tuning a pre-trained model	AdamW	1e-5 – 5e-5	Small LR to not destroy pre-trained weights

Learning rate: the most important hyperparameter

If this book could only teach one hyperparameter lesson, it would be this: the learning rate matters more than the optimizer choice. Too large → loss explodes or oscillates (NaN losses, Chapter 10). Too small → training crawls and stalls in a mediocre spot. The standard practice:

• Start with the table's typical value.
• If loss goes to NaN or explodes in the first epochs → LR too high; divide by 10.
• If loss barely moves after several epochs → LR too low; multiply by 3–10.
• For a principled search, run the LR range test: train briefly while sweeping LR exponentially and pick the value where loss falls fastest (implemented in libraries; conceptually simple).

Learning rate schedulers adjust LR during training — the near-universal pattern is to start higher and decay:

optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
# ... in the epoch loop, after optimizer.step():
scheduler.step()   # every 30 epochs, LR *= 0.1

Common schedulers: StepLR (drop by gamma every step_size epochs), CosineAnnealingLR (smooth cosine decay — popular in modern papers), ReduceLROnPlateau (drop LR when validation loss stalls). You don't need schedulers for this book's toy examples, but expect them in every serious training script you read.

Worked example: the same data, three losses

To make the loss/task matching concrete, here's one dataset trained three ways — the only things changing are the head, targets, and loss:

import torch, torch.nn as nn
torch.manual_seed(0)

X = torch.randn(500, 4)
# Task A: regression target
y_reg = X.sum(dim=1, keepdim=True) + torch.randn(500, 1) * 0.1
# Task B: binary label from the same signal
y_bin = (y_reg > 0).float()
# Task C: 3-class label
y_cls = torch.bucketize(y_reg.squeeze(), torch.tensor([-1.0, 1.0]))  # 0,1,2

def train(head_out, y, loss_fn, epochs=400):
    m = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, head_out))
    opt = torch.optim.AdamW(m.parameters(), lr=1e-2)
    for _ in range(epochs):
        opt.zero_grad(); loss = loss_fn(m(X), y); loss.backward(); opt.step()
    return m, loss.item()

_, l_reg = train(1, y_reg, nn.MSELoss())
_, l_bin = train(1, y_bin, nn.BCEWithLogitsLoss())
_, l_cls = train(3, y_cls.long(), nn.CrossEntropyLoss())
print(f"regression MSE: {l_reg:.4f} | binary BCE: {l_bin:.4f} | 3-class CE: {l_cls:.4f}")

Same features, same network shape, three tasks — only the loss (and target format) changes. That is the lesson: the loss function is how you tell the network what problem it's solving.

For Your Research

Loss choice is a methods-section decision — justify it. Reviewers notice when a paper uses an unusual loss without explanation, and they also notice a well-motivated one. The default expectation: cross-entropy for classification, MSE for regression. If you deviate (focal loss for class imbalance, contrastive losses for representation learning, a custom composite loss), your paper needs one or two sentences saying why — e.g., "we use focal loss because positives are outnumbered 100:1." Keep a lab habit: every time you try a non-default loss in an experiment, write the one-sentence justification in your experiment log (Chapter 11). Half your future paper's methods section will assemble itself from those sentences.

Class imbalance: weighted losses

Real datasets are rarely balanced — imagine 95% negative, 5% positive samples. Plain cross-entropy lets the model get 95% accuracy by always predicting negative, learning nothing. The standard fix is class weights: penalize mistakes on rare classes more:

# class_counts = [950, 50] -> weight rare class 19x more
weights = torch.tensor([1.0, 19.0])
loss_fn = nn.CrossEntropyLoss(weight=weights)

For binary tasks, BCEWithLogitsLoss(pos_weight=torch.tensor([19.0])) does the same. Rule of thumb: set weights inversely proportional to class frequencies, then tune. (More advanced: focal loss, which down-weights easy examples — worth knowing by name for your literature reviews on imbalanced problems.)

Writing a custom loss

Sometimes the standard losses don't express your goal — e.g., you want MSE but with extra penalty on under-prediction. A loss is just a differentiable function returning a scalar; write it as one:

def asymmetric_mse(pred, target, under_penalty=2.0):
    err = pred - target
    # penalize under-prediction (err < 0) more heavily
    weight = torch.where(err < 0, torch.full_like(err, under_penalty), torch.ones_like(err))
    return (weight * err ** 2).mean()

# use it exactly like a built-in loss:
loss = asymmetric_mse(model(xb), yb)
loss.backward()  # autograd handles the rest

Because it's built from tensor ops, autograd differentiates it automatically (Chapter 4's promise in action). Custom losses are a legitimate, common paper contribution — "we propose a loss that..." — and now you know they're just functions.

Optimizer shootout: feel the difference

Theory is nice; watching optimizers behave differently builds intuition no table can. This trains the same tiny problem with three optimizers and prints the trajectory:

import torch, torch.nn as nn
torch.manual_seed(0)
X = torch.randn(200, 4); y = (X.sum(dim=1, keepdim=True) > 0).float()

def run(opt_name, lr):
    m = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 1))
    opt = {"sgd": torch.optim.SGD(m.parameters(), lr=lr),
           "sgd-m": torch.optim.SGD(m.parameters(), lr=lr, momentum=0.9),
           "adamw": torch.optim.AdamW(m.parameters(), lr=lr)}[opt_name]
    losses = []
    for _ in range(100):
        opt.zero_grad()
        l = nn.BCEWithLogitsLoss()(m(X), y)
        l.backward(); opt.step()
        losses.append(l.item())
    return losses[0], losses[9], losses[99]

for name, lr in [("sgd", 0.1), ("sgd-m", 0.1), ("adamw", 0.01)]:
    l0, l10, l100 = run(name, lr)
    print(f"{name:6s} loss: start {l0:.3f} -> ep10 {l10:.3f} -> ep100 {l100:.3f}")

Typical result: AdamW drops fastest early, SGD+momentum catches up, plain SGD lags — the standard story in one screen of output. Run it yourself with different LRs; the feel of "AdamW is forgiving, SGD needs its LR tuned" will stick.

Reduction modes and label smoothing: two refinements worth knowing

Reduction controls how per-sample losses combine: reduction="mean" (default — one scalar for backward), "sum", or "none" (keep per-sample losses). "none" is useful when you need per-sample weighting or analysis:

loss_fn = nn.CrossEntropyLoss(reduction="none")
per_sample = loss_fn(model(xb), yb)   # shape (B,) — inspect, weight, or mask
loss = per_sample.mean()              # then reduce yourself

Researchers use this for importance weighting, hard-example mining, and loss landscapes per sample. If a paper says "we weight each sample's loss by...", this is the mechanism.

Label smoothing (a CrossEntropyLoss argument) softens hard 0/1 targets slightly (e.g., correct class gets 0.9 instead of 1.0), discouraging overconfident predictions:

loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)

It's a one-argument regularizer that helps in many classification settings — cheap to try in a sweep (Chapter 11) when a model overfits.

When MSE misleads

MSE's quadratic penalty assumes errors are roughly symmetric and outlier-free. Two common departures: (1) outliers — a few bad labels dominate the loss; switch to L1Loss or SmoothL1Loss. (2) Heteroscedastic noise — some inputs are inherently noisier than others; consider predicting uncertainty too (a model outputting mean and variance, trained with Gaussian negative log-likelihood — an intermediate topic you'll recognize when you meet it). The lesson generalizes: the loss encodes your assumptions about the noise — choose it as deliberately as the architecture.

Key Takeaways

1. Match the loss to the task: MSE for regression, BCEWithLogitsLoss for binary, CrossEntropyLoss for multi-class.
2. Feed raw logits (no softmax/sigmoid) to BCEWithLogitsLoss and CrossEntropyLoss; targets are floats for BCE, int64 class indices for CE.
3. AdamW (lr≈1e-3) is the default optimizer for new work; SGD + momentum (lr≈0.1) for CNNs trained from scratch; tiny LRs (1e-5) for fine-tuning.
4. The learning rate matters more than the optimizer — too high explodes, too low stalls; schedulers decay it during training.
5. Weight decay (L2) is your baseline regularization; set it (1e-4) unless you have a reason not to.

---

Chapter 7: The Training Loop, Step by Step

What the loop actually does

Strip away every library, and training a neural network is one repeated procedure:

1. Take a batch of data.
2. Run it through the model (forward pass) → predictions.
3. Compare predictions to targets with the loss function → one number.
4. Compute gradients (backward pass).
5. Nudge every weight a small step against its gradient (optimizer step).
6. Repeat until the loss stops improving — then check on data the model never saw (validation).

Every training script ever written — from your first MNIST classifier to a 70-billion-parameter language model — is this loop with better engineering around it. This chapter builds the full, honest version: device handling, mini-batches, train/validation split, metric tracking, and clean printing. Memorize its shape; you will write it hundreds of times.

The pieces, assembled

We'll train a classifier on a synthetic 3-class problem (so it runs instantly on CPU), with a proper validation split. Read it slowly — every line earns its place:

import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader, random_split

torch.manual_seed(0)

# ---------- 1. Data ----------
N = 3000
X = torch.randn(N, 8)
# 3 classes from a nonlinear rule + noise (so a linear model can't solve it)
logits_true = torch.stack([
    X[:, 0] + X[:, 1],
    X[:, 2] - X[:, 3] + 0.5 * X[:, 4] ** 2,
    -X[:, 0] + X[:, 5],
], dim=1)
y = logits_true.argmax(dim=1) + (torch.rand(N) < 0.05).long() * 0  # clean labels
y = logits_true.argmax(dim=1)

full = TensorDataset(X, y)
train_ds, val_ds = random_split(full, [2400, 600],
                                generator=torch.Generator().manual_seed(0))
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=256)

# ---------- 2. Model, loss, optimizer ----------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = nn.Sequential(
    nn.Linear(8, 64), nn.ReLU(),
    nn.Linear(64, 64), nn.ReLU(),
    nn.Linear(64, 3),
).to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2, weight_decay=1e-4)

# ---------- 3. Training loop ----------
epochs = 25
for epoch in range(1, epochs + 1):
    # --- train ---
    model.train()
    train_loss, correct, total = 0.0, 0, 0
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        optimizer.zero_grad()          # 1. clear old gradients
        out = model(xb)                # 2. forward
        loss = loss_fn(out, yb)        # 3. loss
        loss.backward()                # 4. backward
        optimizer.step()               # 5. update weights

        train_loss += loss.item() * xb.size(0)
        correct += (out.argmax(dim=1) == yb).sum().item()
        total += xb.size(0)
    train_loss /= total
    train_acc = correct / total

    # --- validate ---
    model.eval()
    val_loss, correct, total = 0.0, 0, 0
    with torch.no_grad():              # no graphs needed for evaluation
        for xb, yb in val_loader:
            xb, yb = xb.to(device), yb.to(device)
            out = model(xb)
            val_loss += loss_fn(out, yb).item() * xb.size(0)
            correct += (out.argmax(dim=1) == yb).sum().item()
            total += xb.size(0)
    val_loss /= total
    val_acc = correct / total

    print(f"epoch {epoch:2d} | train loss {train_loss:.4f} acc {train_acc:.3f} "
          f"| val loss {val_loss:.4f} acc {val_acc:.3f}")

Run it. You should see train loss falling from ~1.1 toward ~0.1, with validation accuracy climbing into the 0.95+ range. If your numbers differ slightly, that's fine — randomness (Chapter 11 covers taming it).

Reading the loop line by line

Data (section 1). TensorDataset wraps tensors so the loader can index them; random_split carves off validation data with a seeded generator (reproducible splits — note this, Chapter 11); DataLoader yields mini-batches, shuffling the training set each epoch. Why shuffle training but not validation? Shuffling training batches makes each gradient step see a different mix — the "stochastic" in SGD — which helps generalization. Validation just needs a stable, representative measurement.

Why mini-batches at all? Three reasons: (1) the full dataset rarely fits in GPU memory; (2) noisy per-batch gradients actually help escape sharp, poor minima; (3) more updates per epoch = faster learning. Batch size 32–256 is the normal range; it is itself a hyperparameter (larger batches train faster per epoch but may generalize slightly worse).

The train block. model.train() sets training mode (matters for dropout/BatchNorm, Chapter 5). The five-step rhythm from Chapter 4 is right there: zero_grad → forward → loss → backward → step. Metrics accumulate with loss.item() * batch_size — multiplying back by batch size because CrossEntropyLoss returns the mean over the batch, and the last batch may be smaller; weighting by actual batch size keeps the epoch average exact.

The validation block. model.eval() + torch.no_grad(): no dropout randomness, no gradient graphs — faster and less memory. We never call optimizer.step() here; validation data must not influence weights, or it stops being an honest measure. The validation accuracy is the number you actually care about — training accuracy tells you the model can learn; validation accuracy tells you it did learn something general.

How to read the printed curves (without plotting)

epoch  1 | train loss 0.8234 acc 0.612 | val loss 0.6012 acc 0.745
epoch  5 | train loss 0.2145 acc 0.921 | val loss 0.1988 acc 0.933
epoch 15 | train loss 0.0512 acc 0.986 | val loss 0.0891 acc 0.972
epoch 25 | train loss 0.0188 acc 0.997 | val loss 0.0764 acc 0.978

Healthy training looks like this: both losses fall, both accuracies rise, validation tracks training with a small gap. Warning signs to recognize:

• Train loss falls, val loss rises → overfitting (Book 6 of this series; fixes: more data, regularization, early stopping).
• Both stuck high → underfitting: model too small, LR too low, or a bug (Chapter 10).
• Loss = NaN → LR too high or numerical instability (Chapter 10).
• Val much better than train → usually dropout/BatchNorm behaving differently in the two modes — actually normal, not a bug.

A cleaner pattern: functions for train and evaluate

As scripts grow, factor the loop into functions. This is the shape real research code takes:

def train_one_epoch(model, loader, loss_fn, optimizer, device):
    model.train()
    total_loss, correct, total = 0.0, 0, 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        optimizer.zero_grad()
        out = model(xb)
        loss = loss_fn(out, yb)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * xb.size(0)
        correct += (out.argmax(dim=1) == yb).sum().item()
        total += xb.size(0)
    return total_loss / total, correct / total

@torch.no_grad()
def evaluate(model, loader, loss_fn, device):
    model.eval()
    total_loss, correct, total = 0.0, 0, 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        out = model(xb)
        total_loss += loss_fn(out, yb).item() * xb.size(0)
        correct += (out.argmax(dim=1) == yb).sum().item()
        total += xb.size(0)
    return total_loss / total, correct / total

# then the driver is five readable lines per epoch:
for epoch in range(1, epochs + 1):
    tr_loss, tr_acc = train_one_epoch(model, train_loader, loss_fn, optimizer, device)
    va_loss, va_acc = evaluate(model, val_loader, loss_fn, device)
    print(f"epoch {epoch:2d} | train {tr_loss:.4f}/{tr_acc:.3f} | val {va_loss:.4f}/{va_acc:.3f}")

The @torch.no_grad() decorator on evaluate is equivalent to the with block — a tidy idiom you'll see in paper repositories.

For Your Research

Your training loop is the most-reviewed code you will write. When you submit code with a paper (or a reviewer asks for it), the training script is what they run first. Make it boring in the best way: fixed seeds, explicit device handling, separate train/eval functions, logged metrics every epoch, and saved checkpoints (Chapter 9). A reviewer who can run python train.py --config configs/baseline.yaml and reproduce your table's numbers is a reviewer who trusts your results. Every "boring" engineering choice in this chapter — seeded splits, .item() logging, eval mode — is a trust-building choice. Chapter 11 turns this loop into a full experiment harness.

Early stopping: quit while you're ahead

Training longer isn't always better — past a point, the model memorizes training noise and validation worsens (overfitting). Early stopping watches validation loss and halts when it stops improving, keeping the best weights:

best_val = float("inf"); patience = 5; wait = 0
best_state = None

for epoch in range(1, 200):
    tr_loss, tr_acc = train_one_epoch(model, train_loader, loss_fn, optimizer, device)
    va_loss, va_acc = evaluate(model, val_loader, loss_fn, device)

    if va_loss < best_val:
        best_val, wait = va_loss, 0
        best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
    else:
        wait += 1
        if wait >= patience:
            print(f"early stop at epoch {epoch}, best val {best_val:.4f}")
            break

model.load_state_dict(best_state)  # restore the best weights, not the last ones

patience (typically 5–20) is how many stagnant epochs you tolerate — it forgives noisy validation curves. Note the .cpu().clone(): it snapshots weights off-GPU so later training steps don't mutate your saved best. Early stopping is the simplest regularization that actually works, and it belongs in every serious training script.

Gradient clipping: a seatbelt for training

Occasionally gradients explode (very deep nets, RNNs, high LRs) — one bad batch sends weights to NaN and the run is dead. Gradient clipping caps the total gradient norm before the optimizer step:

optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

One line, placed between backward() and step(). It doesn't change healthy training (norms below 1.0 are untouched) but prevents catastrophic steps. Cheap insurance — many researchers add it by default to new projects.

Plotting learning curves

Printed numbers are fine; curves are better. Collect metrics during training and plot them — the shape of the curves diagnoses everything from Chapter 7's "how to read" section at a glance:

import matplotlib.pyplot as plt

history = {"train_loss": [], "val_loss": [], "val_acc": []}
for epoch in range(1, epochs + 1):
    tr_loss, _ = train_one_epoch(model, train_loader, loss_fn, optimizer, device)
    va_loss, va_acc = evaluate(model, val_loader, loss_fn, device)
    history["train_loss"].append(tr_loss)
    history["val_loss"].append(va_loss)
    history["val_acc"].append(va_acc)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
ax1.plot(history["train_loss"], label="train"); ax1.plot(history["val_loss"], label="val")
ax1.set(title="Loss", xlabel="epoch"); ax1.legend()
ax2.plot(history["val_acc"], label="val acc", color="green")
ax2.set(title="Validation accuracy", xlabel="epoch"); ax2.legend()
plt.savefig("learning_curves.png"); plt.show()

Make this plot for every experiment and save the PNG next to metrics.csv (Chapter 11). Six months later, the curve will tell you more in five seconds than the CSV will in five minutes.

Beyond accuracy: the confusion matrix

Accuracy hides which mistakes a model makes. A confusion matrix shows them — rows are true classes, columns are predictions:

@torch.no_grad()
def confusion_matrix(model, loader, device, num_classes):
    model.eval()
    cm = torch.zeros(num_classes, num_classes, dtype=torch.long)
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        pred = model(xb).argmax(dim=1)
        for t, p in zip(yb.cpu(), pred.cpu()):
            cm[t, p] += 1
    return cm

cm = confusion_matrix(model, val_loader, device, 3)
print(cm)
# tensor([[190,   8,   2],
#         [ 12, 175,  13],
#         [  3,  11, 186]])

Read it: row 1 (true class 1) shows 12 samples misclassified as class 0 and 13 as class 2 — if one pair of classes confuses the model systematically, that's a finding: maybe the classes genuinely overlap, maybe the features are insufficient, maybe the labels are noisy. Papers report confusion matrices precisely because they turn "94% accuracy" into an analysis of the remaining 6% — and the remaining errors are often where the next idea lives.

Sanity checks before long runs

Before launching anything over an hour, run this five-minute gauntlet:

1. One-batch overfit (Chapter 10) — can it learn at all?
2. One full epoch — does the pipeline run end-to-end, and what's the samples/second (extrapolate total time)?
3. Checkpoint save/load round-trip (Chapter 9) — save after epoch 1, reload, verify allclose.
4. Overnight small run — train the small config fully; confirm curves look healthy before scaling up.

Each check is cheap; skipping them is how you discover at hour 30 that checkpoints were never written.

The learning-rate range test: stop guessing

Instead of guessing the LR, measure it: train for a few hundred batches while increasing the LR exponentially each batch, and plot loss vs. LR. Loss falls steeply in a "good" LR window, then explodes — pick an LR near the steepest descent, typically ~10× below the explosion point:

lrs, losses = [], []
model = FreshModel().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-7)  # start tiny
lr_mult = 1.02  # grow ~2% per batch

model.train()
for i, (xb, yb) in enumerate(train_loader):
    if i > 300: break
    xb, yb = xb.to(device), yb.to(device)
    opt.zero_grad()
    loss = loss_fn(model(xb), yb)
    loss_val = loss.item()
    if not losses:
        init_loss = loss_val
    if torch.isnan(loss) or loss_val > 4 * init_loss:
        break  # diverged — the window just ended
    loss.backward(); opt.step()
    lrs.append(opt.param_groups[0]["lr"]); losses.append(loss_val)
    opt.param_groups[0]["lr"] *= lr_mult

import matplotlib.pyplot as plt
plt.semilogx(lrs, losses); plt.xlabel("lr (log)"); plt.ylabel("loss")
plt.savefig("lr_range_test.png")

This technique (from Leslie Smith's work on cyclical learning rates) turns the most important hyperparameter from a guess into a five-minute measurement. Run it once per new architecture/dataset pair; the plot goes next to your learning curves in the run directory.

Key Takeaways

1. The training loop is always: batch → forward → loss → backward → step, repeated per epoch, with a validation pass (eval mode, no grad, no optimizer steps) measuring honest progress.
2. Shuffle training batches; keep validation stable and seeded.
3. Weight metric averages by actual batch size (loss.item() * batch_size).
4. Validation accuracy is the number that matters — training metrics measure capacity, validation measures learning.
5. Learn to read the curves: both falling = healthy; train-only falling = overfitting; NaN = LR/numerical trouble.
6. Factor the loop into train_one_epoch / evaluate functions — that is the shape of real research code.

---

Chapter 8: Datasets and DataLoaders — Feeding Real Data

Why not just use tensors directly?

Chapter 7 used TensorDataset — tensors wrapped for the loader. That works when your whole dataset fits in memory as tensors. Real research data rarely cooperates: images live as thousands of JPEG files on disk, text as raw documents, sensor readings as CSVs too large for RAM. You need a system that loads samples on demand, preprocesses them identically every time, and serves them in efficient batches — possibly using multiple background workers so the GPU never waits for data.

That system is Dataset + DataLoader, and understanding it deeply will save you from the most tedious class of research bugs: train/validation preprocessing mismatches, slow data pipelines, and memory blowups.

The Dataset contract: len and getitem

A Dataset is any class with two methods: __len__ (how many samples) and __getitem__(i) (return sample i). That's the entire contract — beautifully minimal:

import torch
from torch.utils.data import Dataset

class NumberDataset(Dataset):
    """Toy dataset: sample i is (features, label) with label = parity of i."""
    def __init__(self, n):
        self.n = n
    def __len__(self):
        return self.n
    def __getitem__(self, i):
        x = torch.tensor([float(i), float(i % 7)], dtype=torch.float32)
        y = torch.tensor(i % 2, dtype=torch.long)
        return x, y

ds = NumberDataset(1000)
print(len(ds))       # 1000
print(ds[42])        # (tensor([42., 0.]), tensor(0))

For real data, __getitem__ loads and preprocesses one sample — reads an image file, tokenizes a sentence, normalizes a signal. Keep __getitem__ deterministic given the index (no randomness unless it's intentional augmentation with a seeded generator): reproducibility (Chapter 11) depends on the data pipeline being stable.

The DataLoader: batching, shuffling, workers

The DataLoader wraps a dataset and yields batches:

from torch.utils.data import DataLoader

loader = DataLoader(ds, batch_size=32, shuffle=True, num_workers=0)
for xb, yb in loader:
    print(xb.shape, yb.shape)  # (32, 2) (32,) — then (32,2)... last batch may be smaller
    break

Key arguments:

• batch_size — samples per batch (32–256 typical).
• shuffle=True for training (reshuffles each epoch), False for validation/test.
• num_workers — background processes that run __getitem__ in parallel. 0 (default) loads in the main process — fine for small/toy data. For image datasets on disk, num_workers=4 (or 8) keeps the GPU fed; the rule of thumb is 2–4 workers per GPU. On Windows, workers require the script guard (if __name__ == "__main__":).
• pin_memory=True — speeds up transfer of batches to CUDA (use with GPU training).
• drop_last=True — drops the final smaller batch; useful when batch-size-sensitive ops (like BatchNorm) misbehave on tiny batches.

How the batch is assembled: the loader calls __getitem__ batch_size times and stacks the results with the default collate_fn. It handles the standard cases (tensors → stacked; numbers → tensored). You only write a custom collate_fn when samples have variable sizes — e.g., sentences of different lengths that need padding to the batch's max length. That's an NLP rite of passage; for now, know it exists.

Worked example: a custom image-like dataset from files

This is the pattern behind most vision research code — samples as files on disk, labels from directory names or a CSV:

import os, csv
import torch
from torch.utils.data import Dataset, DataLoader

class CSVDataset(Dataset):
    """Each row: feature_1,...,feature_8,label. Loaded lazily per sample."""
    def __init__(self, csv_path):
        with open(csv_path) as f:
            rows = list(csv.reader(f))
        self.header, self.rows = rows[0], rows[1:]

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, i):
        *feats, label = self.rows[i]
        x = torch.tensor([float(v) for v in feats], dtype=torch.float32)
        y = torch.tensor(int(label), dtype=torch.long)
        return x, y

# create a demo CSV so the example runs standalone
torch.manual_seed(1)
with open("/tmp/demo.csv", "w", newline="") as f:
    w = csv.writer(f)
    w.writerow([f"f{j}" for j in range(8)] + ["label"])
    for _ in range(500):
        feats = torch.randn(8)
        label = int(feats[:2].sum() > 0)
        w.writerow([f"{v:.4f}" for v in feats] + [label])

ds = CSVDataset("/tmp/demo.csv")
train_ds, val_ds = torch.utils.data.random_split(
    ds, [400, 100], generator=torch.Generator().manual_seed(0))
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
print(len(train_loader))  # 13 batches (12x32 + 1x16)
xb, yb = next(iter(train_loader))
print(xb.shape, yb.shape, yb.dtype)  # (32, 8) (32,) torch.int64

Note the dtype discipline: features float32, labels long — exactly what CrossEntropyLoss expects (Chapter 6). Getting dtypes right in __getitem__ prevents a whole class of downstream errors.

Transforms: torchvision's preprocessing pipeline

For images, torchvision provides datasets (MNIST, CIFAR-10, ...) and transforms — composable preprocessing:

from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),                          # PIL image -> float tensor, scales to [0,1]
    transforms.Normalize((0.5,), (0.5,)),           # -> roughly [-1, 1]
])

train_ds = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=2)
xb, yb = next(iter(train_loader))
print(xb.shape)  # (64, 1, 28, 28): (batch, channels, height, width)

The canonical image tensor layout is (N, C, H, W) — batch, channels, height, width. (Note: this is "channels-first"; TensorFlow/Keras defaults to channels-last. Mixing them up is a classic porting bug.) Normalize(mean, std) uses dataset statistics — for MNIST (0.1307,), (0.3081,) are the standard values; for ImageNet-pretrained models use ImageNet's ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]).

Critical rule: identical preprocessing for train and validation, except augmentation. Data augmentation (random crops, flips) applies to training only; validation gets the deterministic pipeline. The standard structure:

train_tf = transforms.Compose([
    transforms.RandomHorizontalFlip(),   # augmentation: train only
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize(mean, std),
])
val_tf = transforms.Compose([
    transforms.ToTensor(),               # deterministic: val/test
    transforms.Normalize(mean, std),
])

A train/val preprocessing mismatch (e.g., normalizing one but not the other) produces the baffling symptom "great training accuracy, terrible validation accuracy" with no overfitting in sight — Chapter 10 lists it.

Performance: keeping the GPU fed

Data loading is the silent killer of training speed. If your GPU sits idle waiting for batches, wall-clock time explodes. The checklist:

1. num_workers=4 (tune: more isn't always better; watch CPU RAM).
2. pin_memory=True when training on CUDA.
3. Do heavy preprocessing once, offline (resize images, tokenize text, save as tensors) rather than in __getitem__ every epoch.
4. If the dataset fits in RAM, preload it (like our CSV example holding rows in memory) instead of re-reading files.
5. Profile with a quick timing loop before blaming the model:

import time
t0 = time.time(); n = 0
for xb, yb in train_loader:
    n += xb.size(0)
print(f"{n} samples in {time.time()-t0:.1f}s -> {n/(time.time()-t0):.0f} samples/s")

If samples/second is low and GPU utilization is low, your bottleneck is data loading, not the model — fix the pipeline before buying a bigger GPU.

For Your Research

Document your data pipeline like a method. Papers get rejected (or results questioned) over vague data handling: "we normalized the images" — how, with what statistics, in what order? Your Dataset class and transform pipeline are the precise answer, so keep them in version control and reference them: "preprocessing followed the pipeline in data/datasets.py (resize to 256, center-crop 224, normalize with ImageNet statistics)." Better still, log the transform list with each experiment (Chapter 11). When another lab emails "we couldn't reproduce your preprocessing," you send a file, not a paragraph from memory.

Variable-length data: writing a custom collate_fn

The default collate stacks same-shaped samples. For variable-length sequences (sentences, time series), you must pad each batch to its longest sequence. That's a custom collate_fn:

import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence

class SeqDataset(Dataset):
    def __init__(self, n=200):
        torch.manual_seed(0)
        self.seqs = [torch.randn(torch.randint(3, 10, (1,)).item(), 4) for _ in range(n)]
        self.labels = torch.randint(0, 2, (n,))
    def __len__(self): return len(self.seqs)
    def __getitem__(self, i): return self.seqs[i], self.labels[i]

def pad_collate(batch):
    seqs, labels = zip(*batch)
    padded = pad_sequence(seqs, batch_first=True)          # (B, T_max, 4), padded with 0
    lengths = torch.tensor([len(s) for s in seqs])         # remember true lengths!
    return padded, lengths, torch.stack(labels)

loader = DataLoader(SeqDataset(), batch_size=8, collate_fn=pad_collate)
xb, lens, yb = next(iter(loader))
print(xb.shape, lens, yb.shape)  # e.g. (8, 9, 4) tensor([...]) (8,)

Two lessons: (1) pad_sequence does the mechanical padding; (2) always return the true lengths — the model needs them to ignore padding (via pack_padded_sequence for RNNs, or attention masks for transformers). Forgetting the lengths and letting the model "read" the zeros is a silent correctness bug: the model learns patterns in your padding.

Beyond MNIST: the dataset zoo

torchvision.datasets ships far more than MNIST, all with the same interface:

from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
                         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

cifar_train = datasets.CIFAR10(root="./data", train=True, download=True, transform=tf)
cifar_test = datasets.CIFAR10(root="./data", train=False, download=True, transform=tf)
print(len(cifar_train), cifar_train[0][0].shape, cifar_train.classes[:3])
# 50000 torch.Size([3, 32, 32]) ['plane', 'car', 'bird']

CIFAR-10 (60k tiny color images, 10 classes) is the canonical "my first real CNN" dataset — small enough for Colab, hard enough to be interesting. Also worth knowing: FashionMNIST (drop-in MNIST replacement, harder), ImageFolder (builds a dataset from a directory of class-named folders — the standard for custom image data), and Hugging Face datasets (load_dataset("...")) for the modern one-liner experience across thousands of datasets.

ImageFolder deserves special mention because it's how most custom vision projects start:

data/
  train/
    cats/  -> img001.jpg, img002.jpg, ...
    dogs/  -> img101.jpg, ...
  val/
    cats/ ...
    dogs/ ...

train_ds = datasets.ImageFolder("data/train", transform=train_tf)
print(train_ds.classes, train_ds.class_to_idx)  # (['cats', 'dogs'], {'cats': 0, 'dogs': 1})

No CSV, no custom class — folder names become labels. For your first real image project, this is the path of least resistance.

Imbalanced data: weighted sampling

Chapter 6 fixed imbalance with loss weights. The complementary fix is at the data level: sample rare classes more often with WeightedRandomSampler:

from torch.utils.data import WeightedRandomSampler

# class_counts: how many samples per class in the training set
class_counts = torch.tensor([950, 50])
class_weights = 1.0 / class_counts
sample_weights = class_weights[train_labels]  # weight per sample by its class

sampler = WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), replacement=True)
train_loader = DataLoader(train_ds, batch_size=32, sampler=sampler)
# NOTE: sampler and shuffle=True are mutually exclusive — the sampler defines the order

Each batch now sees roughly balanced classes even though the dataset isn't. Use loss weighting or balanced sampling (rarely both — that double-corrects). Which to choose? Sampling changes what the model sees; weighting changes what it cares about. Sampling is often more stable for extreme imbalance; weighting is simpler. Try both in a sweep — that's what sweeps are for.

Subset and random_split: carving datasets precisely

Beyond random_split, Subset wraps any index list — useful for debugging on 100 samples or building custom folds:

from torch.utils.data import Subset
tiny = Subset(full_dataset, indices=list(range(100)))
tiny_loader = DataLoader(tiny, batch_size=32)   # 4 batches — the debugging workhorse

The "debug on a tiny subset" habit (Chapter 10's minimal reproduction) starts here: if your pipeline works on 100 samples, data loading is innocent and the bug is elsewhere.

Squeezing the pipeline: prefetch and persistence

Two DataLoader flags for when you've outgrown the basics:

• persistent_workers=True — keeps worker processes alive between epochs instead of respawning them. With num_workers=4, respawning every epoch can cost seconds per epoch; persistence removes that tax. (Only meaningful with num_workers > 0.)
• prefetch_factor=2 (default) — how many batches each worker prepares ahead. Raise it if the GPU still starves with workers running; lower it if RAM is tight.

And a measurement habit: compare samples/s (Chapter 8's timing loop) before and after each change. Data-pipeline tuning without measurement is superstition — with measurement, it's engineering. One more: if your epoch time increases with more workers, you've hit CPU/RAM contention — back off. The optimum is empirical, usually 2–8 workers.

Three-way splits: train / validation / test, done right

Chapter 7 used train/val. For publishable work you need three splits with strict roles:

from torch.utils.data import random_split

gen = torch.Generator().manual_seed(42)
train_ds, val_ds, test_ds = random_split(full_ds, [0.7, 0.15, 0.15], generator=gen)

• Train — the model learns from this.
• Validation — you tune hyperparameters and pick checkpoints against this. Every decision informed by validation consumes some of its honesty.
• Test — locked away until the very end; run once for the final number. If you tune based on test results, it's validation data wearing a disguise (Chapter 10, Silent bug E).

Practical notes: use the same seeded split for every experiment you compare (or comparisons are meaningless); for small datasets prefer k-fold cross-validation over a single split; and never, ever let test samples leak into training transforms statistics (fit scalers/normalizers on train only).

Key Takeaways

1. Dataset needs only __len__ and __getitem__; DataLoader adds batching, shuffling, and parallel workers.
2. Keep __getitem__ deterministic per index; get dtypes right there (features float32, class labels long).
3. Image tensors are (N, C, H, W) — channels first; normalization statistics must match between train and val.
4. Augmentation is for training only; validation preprocessing must be deterministic.
5. num_workers=4 + pin_memory=True (on CUDA) keeps the GPU fed; profile samples/second before optimizing the model.
6. Your data pipeline is part of your method — version-control it and log it with experiments.

---

Chapter 9: Saving, Loading, and Checkpointing Models

Why saving is a research skill, not housekeeping

Training runs take hours, days, or weeks. GPUs get preempted, Colab runtimes disconnect, laptops run out of battery. Every hour of training that isn't checkpointed is an hour you may have to redo — and in research, "redo" also means re-verifying that the rerun matches, which it won't exactly (Chapter 11). Worse: a published result whose model weights are lost can't be shared, inspected, or built upon.

PyTorch's saving system has one core idea: save the state_dict, not the model object. A state_dict is an ordinary Python dictionary mapping layer names to tensors — the model's learned numbers, nothing else. It's portable, inspectable, and version-proof in a way that pickled model objects are not.

The essential trio: save, load, verify

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(8, 64), nn.ReLU(), nn.Linear(64, 3))

# SAVE: the recommended way — state_dict only
torch.save(model.state_dict(), "model_weights.pt")

# LOAD: rebuild the architecture first, then load the numbers into it
model2 = nn.Sequential(nn.Linear(8, 64), nn.ReLU(), nn.Linear(64, 3))
model2.load_state_dict(torch.load("model_weights.pt", weights_only=True))
model2.eval()

# VERIFY: same weights -> same outputs
x = torch.randn(4, 8)
print(torch.allclose(model(x), model2(x)))  # True

Three details that matter:

1. torch.save / torch.load use Python's pickle under the hood. The weights_only=True argument (default in recent versions) restricts loading to tensors — a security practice: never torch.load a file from an untrusted source without it, since pickle can execute arbitrary code.
2. You must recreate the architecture before loading. The state_dict is just numbers; the nn.Module gives them meaning. This is why keeping model code in version control alongside weights is non-negotiable.
3. Verify after loading. The allclose check takes one line and catches architecture mismatches, partial loads, and device mixups before they corrupt an experiment.

What NOT to do: torch.save(model, "model.pt") — saving the entire model object. It pickles the class definition path, so loading breaks if you rename the file, move the class, or change PyTorch versions. Tutorials show it because it's short; research code avoids it because it's fragile. The one exception is quick personal checkpoints you will delete within the hour.

Checkpoints: saving the full training state

A state_dict saves the model. A checkpoint saves everything needed to resume training exactly: model weights, optimizer state (Adam's running averages — without them, resuming restarts the optimizer's memory from scratch), the epoch number, and the scheduler state:

def save_checkpoint(path, model, optimizer, scheduler, epoch, best_val, rng_state=None):
    torch.save({
        "epoch": epoch,
        "model_state": model.state_dict(),
        "optimizer_state": optimizer.state_dict(),
        "scheduler_state": scheduler.state_dict() if scheduler else None,
        "best_val": best_val,
    }, path)

def load_checkpoint(path, model, optimizer, scheduler, device):
    ckpt = torch.load(path, map_location=device, weights_only=False)
    model.load_state_dict(ckpt["model_state"])
    optimizer.load_state_dict(ckpt["optimizer_state"])
    if scheduler and ckpt["scheduler_state"]:
        scheduler.load_state_dict(ckpt["scheduler_state"])
    return ckpt["epoch"], ckpt["best_val"]

(weights_only=False is needed here because the checkpoint contains non-tensor metadata like the epoch int. Only load checkpoints you created yourself.)

map_location solves the classic "trained on GPU, loading on CPU" crash: torch.load(path, map_location="cpu") remaps all tensors to CPU regardless of where they were saved. In Colab-to-laptop workflows, you will use this constantly.

The two checkpoint files every experiment keeps

In practice, maintain two files per experiment:

best_val = float("inf")
for epoch in range(1, epochs + 1):
    tr_loss, tr_acc = train_one_epoch(model, train_loader, loss_fn, optimizer, device)
    va_loss, va_acc = evaluate(model, val_loader, loss_fn, device)

    # 1. rolling checkpoint: always the latest state (resume after crashes)
    save_checkpoint("checkpoints/last.pt", model, optimizer, scheduler, epoch, best_val)

    # 2. best-model snapshot: the weights you'd actually publish
    if va_loss < best_val:
        best_val = va_loss
        torch.save(model.state_dict(), "checkpoints/best_model.pt")
        print(f"  -> new best val loss {best_val:.4f}, saved")

• last.pt — overwritten every epoch; if the run dies at epoch 47, you resume from 47, not 0.
• best_model.pt — the weights behind your reported numbers. This is the file you share, analyze, and load for the paper's figures.

Why track "best" by validation loss, not training loss? Training loss keeps falling as the model memorizes; validation loss tells you when generalization peaked. Reporting the final epoch's weights instead of the best epoch's is a quiet, common way to under-report your own results by a point or two.

Resuming: the crash-recovery drill

start_epoch, best_val = 1, float("inf")
if os.path.exists("checkpoints/last.pt"):
    start_epoch, best_val = load_checkpoint(
        "checkpoints/last.pt", model, optimizer, scheduler, device)
    start_epoch += 1   # continue with the NEXT epoch
    print(f"resumed from epoch {start_epoch}")

for epoch in range(start_epoch, epochs + 1):
    ...

On Colab, where runtimes disconnect without warning, this pattern is the difference between a dead experiment and a paused one. Save last.pt to Google Drive (not the ephemeral runtime disk) if the run matters.

Inspecting and surgically editing weights

Because a state_dict is just a dict of tensors, you can inspect and modify it — useful for analysis (and for "For Your Research" below):

sd = model.state_dict()
for name, tensor in sd.items():
    print(f"{name:20s} {str(tuple(tensor.shape)):12s} mean={tensor.mean():+.4f} std={tensor.std():.4f}")

# Example: freeze a layer by zeroing its gradients is done via requires_grad,
# but you can also directly inspect what the model learned:
w = sd["0.weight"]  # first Linear layer's weight matrix
print("dead ReLU check — fraction of near-zero rows:", (w.abs().sum(dim=1) < 1e-6).float().mean().item())

For Your Research

Checkpoints are your paper's evidence locker. Reviewers increasingly ask for trained weights ("please release model checkpoints"), and venues like NeurIPS encourage it. Adopt this convention from your first serious experiment: every reported number in a table must trace to a specific best_model.pt file, and every such file lives next to a log of the exact config and code version that produced it (Chapter 11 ties these together). When you write "our model achieves 94.2%," you should be able to point at the file and the log line. This sounds bureaucratic; in practice it has saved countless researchers during rebuttal, when a reviewer asks for one more ablation and you need the exact weights, not an approximate rerun.

Transfer learning in full: from pre-trained weights to your task

Chapter 5 showed freezing; here's the complete transfer-learning workflow — the single highest-value technique for a student with limited compute, since it converts weeks of someone else's GPU time into your starting point:

import torch.nn as nn
from torchvision.models import resnet18, ResNet18_Weights

# 1. Load pre-trained backbone
weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)
preprocess = weights.transforms()   # the EXACT preprocessing it was trained with — use it!

# 2. Replace the head for your task (e.g., 10 classes)
model.fc = nn.Linear(model.fc.in_features, 10)

# 3. Two-phase training: head first, then everything
for p in model.parameters(): p.requires_grad = False
for p in model.fc.parameters(): p.requires_grad = True
opt = torch.optim.AdamW(model.fc.parameters(), lr=1e-3)
# ... train a few epochs (head adapts to new task) ...

for p in model.parameters(): p.requires_grad = True   # unfreeze all
opt = torch.optim.AdamW(model.parameters(), lr=1e-5)  # tiny LR: don't destroy features
# ... fine-tune a few more epochs ...

Three rules that make or break transfer learning: (1) use the pre-trained model's own preprocessing (weights.transforms()) — mismatched normalization silently cripples accuracy; (2) train the head first with the backbone frozen, so random head weights don't send destructive gradients into good features; (3) fine-tune with a tiny LR (1e-5-ish) — you're polishing, not rebuilding.

Surgery: partial loading with strict=False

Sometimes architectures don't match exactly — you're loading old weights into a modified model (extra layer, different head). strict=False loads whatever matches and reports the rest:

# old checkpoint has layers: fc1, fc2. New model has fc1, fc2, fc3.
missing, unexpected = model.load_state_dict(torch.load("old.pt", weights_only=True),
                                            strict=False)
print("missing:", missing)        # ['fc3.weight', 'fc3.bias'] — randomly initialized, expected
print("unexpected:", unexpected) # [] — nothing in the file was unusable

This is also how you inspect what transferred: the missing/unexpected lists are a precise audit. Researchers use this for ablations like "initialize from checkpoint but re-randomize the top layers" — a clean way to ask which layers hold the useful knowledge.

Exporting for inference: a first look at TorchScript

When a model leaves the research script — a demo, a deployment — you want it runnable without your training code. TorchScript compiles a module into a serializable graph:

model.eval()
example = torch.randn(1, 3, 224, 224)
traced = torch.jit.trace(model, example)   # record the graph for this input shape
traced.save("model_traced.pt")
# later, anywhere: loaded = torch.jit.load("model_traced.pt"); loaded(new_input)

Caveat: tracing records one execution path — data-dependent if statements get frozen to whichever branch ran. For models with dynamic control flow, torch.jit.script compiles the actual code instead. Full deployment (ONNX, mobile) is Chapter 12 territory; this is enough to know the door exists.

Organizing weights: naming, model cards, and sharing

A weights file named model.pt in a folder of twelve model(1).pt files is a reproducibility failure waiting to happen. Adopt a naming convention that encodes provenance:

checkpoints/
  20261008_baseline_lr1e-2_seed42_best.pt
  20261008_baseline_lr1e-2_seed42_last.pt

Date, config nickname, key hyperparameters, seed, and role (best/last) — the filename alone answers "what is this?" And when you share weights (increasingly expected with papers), include a model card: a short README stating the architecture, training data, intended use, known limitations, and the exact config/commit that produced it. Hugging Face's hub made model cards standard practice; even a five-line card in your repo is enormously more useful than a bare .pt file.

Loading across PyTorch versions: what actually breaks

A common worry: "will my checkpoint load in next year's PyTorch?" In practice, state_dicts are remarkably stable — they're just named tensors. What can break: (1) loading a full pickled model object (another reason to save state_dicts, Chapter 9's core rule); (2) custom classes referenced by old checkpoints; (3) very old checkpoints predating weights_only defaults. Mitigation is the same trio: save state_dicts, keep model code in version control, and note the PyTorch version in summary.json (Chapter 11). With those, decade-old checkpoints routinely load fine.

Inference-only sharing: what to ship

When someone just needs to use your model (not continue training), ship the smallest sufficient artifact:

# Export: weights + the metadata needed to use them
torch.save({
    "state_dict": model.state_dict(),
    "config": cfg,                    # architecture hyperparameters
    "classes": train_ds.classes,      # label meanings!
    "preprocessing": "weights.transforms() or explicit params",
    "torch_version": torch.__version__,
}, "released_model.pt")

The classes entry is the one everyone forgets: a checkpoint that outputs class index 2 is useless without knowing index 2 means "cat." Similarly, preprocessing parameters (normalization means/stds, image size) must travel with the weights — a model fed wrongly-normalized inputs fails silently (Chapter 10, Silent bug B). A self-describing artifact like this is what turns "here are my weights" into something another researcher can actually run.

For maximum compatibility, also consider exporting to ONNX (torch.onnx.export) — an open format runnable outside PyTorch entirely. It's the lingua franca for deployment, and reviewers in applied venues appreciate it.

Checkpoint hygiene: size, cleanup, and what to keep

Full checkpoints (model + optimizer + scheduler) for large models get big — Adam's state triples the weight size, so a 1 GB model means ~4 GB per checkpoint. Manage it:

import os, glob

def prune_checkpoints(ckpt_dir, keep_last_n=3):
    """Keep the N most recent last_*.pt; never touch best_model.pt."""
    files = sorted(glob.glob(f"{ckpt_dir}/last_epoch*.pt"),
                   key=os.path.getmtime)
    for f in files[:-keep_last_n]:
        os.remove(f)

Conventions that scale: save rolling checkpoints as last_epoch{N}.pt and prune to the last 2–3; never auto-delete best_model.pt; for very long runs, also keep the checkpoint from the best few epochs, not just one (validation noise means "best" is uncertain). On Colab/Drive, watch your storage quota — a run that dies from a full disk at 3 a.m. is a rite of passage you only need once.

One more consideration: checkpoint frequency. Every epoch is right for short runs; for multi-day runs, save every N steps within epochs too (step-based rather than epoch-based checkpointing) so a crash never costs more than ~30 minutes of compute.

Key Takeaways

1. Save state_dict (the weights dict), not the model object — portable, inspectable, robust.
2. A checkpoint = model state + optimizer state + scheduler state + epoch; it's what lets training resume exactly.
3. Keep two files: last.pt (rolling, for crash recovery) and best_model.pt (best validation, for reporting/sharing).
4. Use map_location when loading across devices; use weights_only=True for untrusted files.
5. Recreate the architecture before loading, and verify with allclose — one line that catches silent corruption.
6. Every reported number must trace to a specific weights file.

---

Chapter 10: Debugging — The 10 Most Common PyTorch Errors and Fixes

Debugging is the job

Here is a secret senior researchers rarely advertise: a large fraction of "research" is debugging. Models don't converge, shapes don't match, GPUs run out of memory, and the error message points at line 200 when the bug is on line 40. The researchers who progress fastest aren't the ones who write bug-free code — they're the ones who recognize failure patterns instantly.

This chapter catalogs the ten errors you will hit, in roughly the order you'll meet them. For each: the symptom, why it happens, and the fix. Keep this chapter bookmarked.

Error 1: Shape mismatch in a layer

Symptom:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x10 and 20x64)

Why: A Linear layer (or @) received the wrong-sized input. Here a batch of shape (32, 10) hit a layer expecting 20 input features.

Fix: Print shapes at every stage. The error message tells you exactly what mismatched — read it literally: "32x10 multiplied by 20x64" means the layer's in_features is 20 but your data has 10. Either fix the data pipeline or the layer definition. Defensive habit: after building a model, run one batch through it immediately and print shapes:

xb, _ = next(iter(train_loader))
print("input:", xb.shape)
out = model(xb)
print("output:", out.shape)  # does this match what the loss expects?

Error 2: Tensor on the wrong device

Symptom:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Why: The model is on GPU but the batch is on CPU (or vice versa). The single most common crash in GPU training.

Fix: The Chapter 2/7 pattern — one device variable, .to(device) for model and every batch:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# inside the loop:
xb, yb = xb.to(device), yb.to(device)

If the error persists, some tensor was created inside the model or loss without going through .to(device) — e.g., a mask or target created with torch.zeros(...) defaults to CPU. Create it on the right device: torch.zeros(..., device=xb.device).

Error 3: Forgetting optimizer.zero_grad()

Symptom: No crash — worse. Loss decreases oddly, oscillates, or explodes; results don't match any tutorial.

Why: Gradients accumulate across batches (Chapter 4). Without zeroing, each step uses the sum of all past gradients — training on a phantom of old data.

Fix: Make the five-step rhythm non-negotiable: zero_grad → forward → loss → backward → step, in that order, every iteration. If loss behaves strangely and there's no error message, check for the missing zero_grad() first — it's the prime suspect.

Error 4: Calling .backward() on a non-scalar

Symptom:

RuntimeError: grad can be implicitly created only for scalar outputs

Why: You called .backward() on a vector/matrix — e.g., loss = loss_fn(out, yb) where you forgot the loss reduces, or you called backward on raw outputs.

Fix: Ensure the thing you call .backward() on is a scalar: loss.backward() where loss = loss_fn(...) with a default reduction ('mean'). If you genuinely need vector gradients, pass gradient=torch.ones_like(output).

Error 5: Sigmoid/softmax before the loss (double-squashing)

Symptom: No crash. Training is mysteriously slow or stuck; accuracy far below expectations.

Why: CrossEntropyLoss and BCEWithLogitsLoss apply softmax/sigmoid internally in a numerically stable form. Applying it yourself first squashes the signal twice — gradients nearly vanish.

Fix:

# WRONG:
loss = nn.CrossEntropyLoss()(torch.softmax(model(xb), dim=1), yb)
# RIGHT:
loss = nn.CrossEntropyLoss()(model(xb), yb)   # raw logits in

Apply softmax/sigmoid only when you need probabilities out (evaluation, inference) — never before these losses.

Error 6: Wrong target dtype/shape for the loss

Symptom:

RuntimeError: Expected target size (32,), got torch.Size([32, 1])
# or
RuntimeError: expected scalar type Long but found Float

Why: CrossEntropyLoss wants class indices as int64 shape (N,); BCEWithLogitsLoss/MSELoss want floats matching the output shape.

Fix: The Chapter 6 table, enforced at the data source:

# classification: squeeze to (N,) and cast
yb = yb.long().squeeze()
# binary/regression: match output shape and float
yb = yb.float()

Better: get dtypes right in Dataset.__getitem__ (Chapter 8) so the fix is in one place.

Error 7: CUDA out of memory

Symptom:

RuntimeError: CUDA out of memory. Tried to allocate 1.2 GiB...

Why: The model + a batch + gradients + optimizer state exceed GPU RAM. Note the optimizer state: Adam stores two extra numbers per parameter — it roughly triples memory vs. the weights alone.

Fixes, in order: 1. Smaller batch size (halve it; e.g., 128 → 64). 2. Smaller model or shorter sequences. 3. Gradient accumulation: simulate a big batch with several small ones —

accum_steps = 4
optimizer.zero_grad()
for i, (xb, yb) in enumerate(train_loader):
    loss = loss_fn(model(xb.to(device)), yb.to(device)) / accum_steps
    loss.backward()                      # accumulates
    if (i + 1) % accum_steps == 0:
        optimizer.step()                 # one update per 4 batches
        optimizer.zero_grad()

1. Free what you don't need: wrap evaluation in torch.no_grad(), del large temporaries, torch.cuda.empty_cache() (rarely the real fix, but harmless).
2. Last resort: mixed precision (torch.cuda.amp) or a bigger GPU (Chapter 12).

Error 8: In-place operation breaking autograd

Symptom:

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation

Why: You modified a tensor in-place (e.g., x += 1, relu_(), or x[mask] = 0) after it was recorded in the graph; autograd needs the original value for the backward pass.

Fix: Use out-of-place ops: x = x + 1 instead of x += 1; nn.ReLU() instead of nn.ReLU(inplace=True) while debugging (in-place ReLU is a memory optimization — fine once the code works, risky while developing). The same applies to weight updates: always inside torch.no_grad() (Chapter 4).

Error 9: Forgetting model.eval() / torch.no_grad() at test time

Symptom: No crash. Test accuracy is noisy, lower than expected, or varies run to run; sometimes it's higher than train accuracy in confusing ways.

Why: Dropout still dropping and BatchNorm using batch statistics during evaluation — the model is being measured with training-time randomness.

Fix: The evaluation ritual, every time:

model.eval()
with torch.no_grad():
    ...evaluate...
model.train()  # switch back if training continues

Error 10: DataLoader crashes and silent slowdowns

Symptoms: RuntimeError: DataLoader worker exited unexpectedly / training absurdly slow / TypeError inside the loop.

Why & fixes: - Worker crashes (num_workers > 0): often shared-memory limits in Docker/Colab or unpicklable objects in Dataset. First debug step: set num_workers=0 — if it works, the bug is in multiprocessing, not your data. - Slow training, idle GPU: num_workers too low or heavy per-sample preprocessing (Chapter 8's checklist). - default_collate errors: samples returning inconsistent types/shapes (e.g., None for a corrupt file). Make __getitem__ total: handle or skip bad samples explicitly, never return None.

The universal debugging protocol

When none of the ten match, run this sequence — it resolves the vast majority of remaining issues:

1. Reproduce minimally. Cut the script to the smallest version that still fails (fewer layers, tiny data, CPU). Most bugs become obvious at small scale.
2. Print shapes, dtypes, devices. Insert print(x.shape, x.dtype, x.device) at each stage. The bug is wherever reality first diverges from expectation.
3. Check the five-step rhythm. zero_grad → forward → loss → backward → step — in order, every iteration.
4. Overfit one batch. The single best training sanity check: train on one batch (e.g., 32 samples) — loss should plummet to ~0 within tens of steps. If it can't memorize 32 samples, something is structurally wrong (bug, not hyperparameters). Only after the one-batch test passes should you tune anything.
5. Binary-search the change. If it worked yesterday, git diff — the bug is in the diff. (Another reason Chapter 11's version-control habit matters.)

For Your Research

The one-batch overfit test belongs in your lab's onboarding. Every new model, every new dataset, every ported paper implementation: before any hyperparameter tuning, before any multi-hour run, overfit a single batch. It takes two minutes and answers the only question that matters first — "is this code capable of learning at all?" Researchers who skip it spend days tuning learning rates on broken code; researchers who run it debug in minutes. When you mentor junior students, teach this test alongside the framework. It's the closest thing deep learning has to a unit test for "my training works."

Silent bugs: when there is no error message

The ten errors above at least crash. The truly expensive bugs run fine and give you wrong conclusions. Watch for these four:

Silent bug A: Data leakage. Validation/test data influences training — e.g., computing normalization statistics over the full dataset (including test), or augmenting before the split. Symptom: suspiciously good validation numbers that collapse on truly new data. Fix: split first, then fit every preprocessing statistic on the training split only.

Silent bug B: Train/val preprocessing mismatch. Training normalizes with one mean/std, validation with another (or none). Symptom: good training accuracy, terrible validation, no overfitting pattern. Fix: define transforms once (Chapter 8's train_tf/val_tf sharing the same Normalize), and diff them when this symptom appears.

Silent bug C: Unseeded comparisons. You compare "model A vs model B" but each ran with different random initialization, shuffling, and splits. Symptom: conclusions that flip when rerun. Fix: same seed and same data split for compared runs (Chapter 11); better, run each config over 3 seeds and report mean ± std — the standard reviewers expect.

Silent bug D: The metric lies. Accuracy on a 99-to-1 imbalanced dataset is 99% for a model that predicts "no" always. Symptom: a number that looks great and means nothing. Fix: match the metric to the task — precision/recall/F1 or balanced accuracy for imbalance, MSE/MAE for regression — and always look at a confusion matrix before celebrating.

A fifth, worth naming: tuning on the test set. Every time you adjust hyperparameters based on test numbers, the test set becomes validation data and your reported numbers become optimistic. Keep a held-out test set you touch once, at the very end. This is methodology, not PyTorch — but PyTorch makes it so easy to re-run that the temptation is constant.

Reading a traceback: a 30-second skill

PyTorch tracebacks look long but have a simple structure: read bottom-up. The last lines name the error and the failing operation; scrolling up finds the first frame in your code (not library internals) — that's where your bug lives. Example:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x10 and 20x64)
  ...
  File "train.py", line 42, in train_one_epoch
    out = model(xb)

Line 42, your code, shape mismatch — diagnosed in seconds. Train yourself to skip library frames on sight; the answer is always in your frames.

The NaN detective: finding where training dies

NaN loss is the most dreaded symptom because it tells you nothing about where it started. The systematic hunt:

1. Detect early. Add a NaN check inside the loop so you catch the exact step instead of discovering it at epoch's end:

if torch.isnan(loss):
    print(f"NaN at epoch {epoch}, batch {i}")
    break

1. Bisect the pipeline. Run the forward pass on the failing batch without training: is the NaN in the data (torch.isnan(xb).any() — corrupt inputs happen), in the model output, or only after the loss? Each answer halves the search space.
2. Usual suspects, in order: learning rate too high (÷10 and retry); log(0) or division by a zero-ish value in a custom loss (add small epsilon: torch.log(x + 1e-8)); mixed-precision overflow (Chapter 12); exploding gradients in deep/recurrent nets (gradient clipping, Chapter 7).
3. Check the data first, always. A surprising fraction of NaN hunts end at a corrupt sample — an image file that's actually HTML, a CSV row with text in a numeric column. Validate inputs before blaming the model.

Determinism debugging: when the bug won't sit still

The worst bugs are intermittent — a crash every ~20 runs. Strategy: make it deterministic first, then debug. Set all seeds (Chapter 11), num_workers=0, cudnn.deterministic=True, and rerun until it reproduces reliably. A bug that reproduces every time is a bug you can bisect; a bug that reproduces "sometimes" is a bug you'll chase for weeks. Only after fixing it do you relax determinism for speed. This "deterministic cage" technique is standard practice in serious debugging — randomness is the enemy of diagnosis.

Key Takeaways

1. Read error messages literally — PyTorch's usually name the exact mismatch (shapes, devices, dtypes).
2. The usual suspects, in order: shapes → devices → dtypes → missing zero_grad() → double-squashed loss → eval mode.
3. OOM? Halve the batch size first; use gradient accumulation to recover large effective batches.
4. Never modify tracked tensors in place; never torch.load untrusted files without weights_only=True.
5. The one-batch overfit test is your universal "does training work at all" check — run it before tuning anything.
6. Debug minimally (small data, CPU, num_workers=0), then scale back up.

---

Chapter 11: From Script to Experiment — Organizing Research Code

Why "it ran on my machine" is not research

So far this book's code has lived in single scripts and notebooks. That is fine for learning — and notebooks are genuinely good for exploration. But research has a stricter standard: an experiment is a claim about what happens under specific conditions, and it only counts if someone else (or you, in six months) can recreate those conditions exactly and get the same result. A notebook with cells run in unknown order, an unrecorded learning rate, and weights saved as model_final_FINAL2.pt does not meet that standard.

This chapter turns a training script into an experiment harness: seeded randomness, configuration files, structured logging, and a directory layout that makes every run reproducible and every result traceable. None of this is glamorous. All of it is what separates "I tried something" from "I ran an experiment."

Seeds: taming randomness (as much as possible)

Deep learning is full of randomness: weight initialization, data shuffling, augmentation, dropout. Two runs with different seeds give slightly different numbers — normal and expected. But unrecorded randomness is a reproducibility killer. The fix is a seed function called once at startup:

import random
import numpy as np
import torch

def set_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # Deterministic algorithms where possible (may cost some speed):
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)

Honest caveats a researcher should know: - Seeding makes runs reproducible on the same machine, same software versions, same GPU. Across different GPUs or PyTorch versions, some operations are nondeterministic at the hardware level — results may differ in the last decimal places. - cudnn.deterministic = True can slow training; many labs seed everything except this for long runs, and enable it only when exact reproducibility is needed (e.g., the final reported run). - Always log the seed with the results. "Seed 42" is part of the method.

Also seed your data splits (Chapter 7's random_split with a seeded generator) — otherwise your "train/val split" differs every run and comparisons between experiments are meaningless.

Configs: no magic numbers in code

Every number that isn't architecture — learning rate, batch size, epochs, seed, data paths — belongs in a config, not hard-coded. The simplest honest config is a YAML file (or a dataclass); the point is that the code never changes between experiments, only the config does:

# config.yaml
seed: 42
data:
  path: ./data
  batch_size: 64
model:
  hidden_dim: 64
  num_layers: 2
  dropout: 0.1
train:
  epochs: 25
  lr: 0.01
  weight_decay: 1e-4
  optimizer: adamw

# train.py (sketch)
import yaml, argparse

parser = argparse.ArgumentParser()
parser.add_argument("--config", required=True)
args = parser.parse_args()
with open(args.config) as f:
    cfg = yaml.safe_load(f)

set_seed(cfg["seed"])
model = build_model(cfg["model"])          # hyperparams from config, not literals
optimizer = torch.optim.AdamW(model.parameters(), lr=cfg["train"]["lr"],
                              weight_decay=cfg["train"]["weight_decay"])

Now an experiment is fully described by python train.py --config configs/exp07_lr001.yaml — a command you can paste into a paper's reproducibility note, a lab wiki, or a message to a collaborator. Keep every config you ever run in configs/ under version control: the history of configs is the history of your experimental decisions.

Logging: the experiment diary

Print statements vanish. A proper experiment writes its results to files: per-epoch metrics (for curves), final metrics (for tables), and the config (for provenance). The minimal version is a CSV plus a JSON summary:

import csv, json, os, time

run_dir = f"runs/{time.strftime('%Y%m%d_%H%M%S')}_lr{cfg['train']['lr']}"
os.makedirs(run_dir, exist_ok=True)

# 1. save the config with the run
with open(f"{run_dir}/config.json", "w") as f:
    json.dump(cfg, f, indent=2)

# 2. per-epoch metrics
with open(f"{run_dir}/metrics.csv", "w", newline="") as f:
    writer = csv.writer(f)
    writer.writerow(["epoch", "train_loss", "train_acc", "val_loss", "val_acc"])
    for epoch in range(1, epochs + 1):
        tr_loss, tr_acc = train_one_epoch(...)
        va_loss, va_acc = evaluate(...)
        writer.writerow([epoch, f"{tr_loss:.4f}", f"{tr_acc:.4f}",
                         f"{va_loss:.4f}", f"{va_acc:.4f}"])

# 3. final summary (what goes into the paper's table)
summary = {"best_val_loss": best_val, "best_val_acc": best_acc,
           "best_epoch": best_epoch, "seed": cfg["seed"],
           "torch_version": torch.__version__}
with open(f"{run_dir}/summary.json", "w") as f:
    json.dump(summary, f, indent=2)

This is deliberately low-tech: no special tools, readable forever, diffable in git. When you're ready for more, the standard upgrades are TensorBoard (torch.utils.tensorboard) for live curves and Weights & Biases for cloud-tracked sweeps — both read the same per-epoch numbers; the discipline of logging them is the part that matters.

One more logging rule: log what you will put in the table, at the time you compute it. Reconstructing "what was the best validation accuracy of the March run?" from memory is how errors enter papers. The summary.json habit makes the paper's results table a copy-paste from files.

Directory layout: a home for everything

A layout that scales from first project to thesis:

my-project/
  configs/            # one YAML per experiment — version controlled
  data/               # raw data (or scripts to fetch it); never edit by hand
  src/
    model.py          # nn.Module definitions
    data.py           # Dataset classes, transforms, loaders
    train.py          # training loop + checkpointing (--config entry point)
    evaluate.py       # loads best_model.pt, reports test metrics
  runs/               # one timestamped dir per run: config.json, metrics.csv,
                      # summary.json, checkpoints/  (git-ignored; big files)
  notebooks/          # exploration only — nothing here counts as a result
  requirements.txt    # or environment.yml — the recorded environment

Rules of thumb: src/ is imported, never copy-pasted; runs/ is written by code, never by hand; notebooks/ are for thinking, src/ is for claiming. When a collaborator (or reviewer) asks "how did you get Table 2?", the answer is: the config in configs/, the code in src/, the outputs in runs/<timestamp>/.

Version control: the lab notebook of code

Use git from day one — not as ceremony, but because the diff is the record of what changed between experiments:

git init
git add src configs requirements.txt
git commit -m "baseline MLP, lr=0.01"
# ... run experiment, get results ...
git tag exp07-baseline   # mark the exact code behind a reported number

The .gitignore should exclude runs/ (large, regenerable) and data/ (large, fetched) — but never exclude configs/. When Chapter 10's debugging protocol says "the bug is in the diff," this is what it means. And when a reviewer asks for an ablation six months later, git checkout exp07-baseline puts you back in the exact room where the result was born.

Putting it together: the experiment checklist

Before launching any run that might end up in a paper, verify:

• [ ] Seed set and logged (set_seed, seed in config and summary)
• [ ] Config file saved with the run (no hard-coded hyperparameters)
• [ ] Data split seeded and logged (same split across compared experiments)
• [ ] Environment recorded (torch.__version__, GPU, requirements.txt)
• [ ] Metrics logged per epoch to metrics.csv; final numbers to summary.json
• [ ] Checkpoints: last.pt rolling + best_model.pt by validation (Chapter 9)
• [ ] Code committed and tagged in git
• [ ] The one-batch overfit test passed (Chapter 10)

This checklist takes five minutes and prevents the two classic disasters: "I can't reproduce my own best result" and "I don't know which run produced Table 3."

For Your Research

Reproducibility is a competitive advantage, not overhead. The ML community is moving steadily toward requiring it — reproducibility checklists at major conferences, artifact evaluation, "papers with code" as the norm. Students who build these habits early produce work that gets cited and built upon, because other researchers can actually use it. More selfishly: the best person these habits serve is you, three months from now, trying to extend your own paper for the journal version. Future-you will thank present-you for every config file, every seed, every summary.json. Treat experiment organization as part of the research contribution, not as janitorial work before the "real" research starts.

TensorBoard in fifteen lines

CSV logging is forever-readable; TensorBoard adds live curves you can watch during training. It's built into PyTorch — no new concepts, just a writer object:

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(f"runs/tb_lr{cfg['train']['lr']}")  # one dir per run
for epoch in range(1, epochs + 1):
    tr_loss, tr_acc = train_one_epoch(...)
    va_loss, va_acc = evaluate(...)
    writer.add_scalar("loss/train", tr_loss, epoch)
    writer.add_scalar("loss/val", va_loss, epoch)
    writer.add_scalar("acc/val", va_acc, epoch)
    writer.add_scalar("lr", optimizer.param_groups[0]["lr"], epoch)
writer.close()
# then: tensorboard --logdir runs  (open the URL it prints)

Log the learning rate too — when a run behaves oddly, the first question is "what was the LR doing?", and schedulers make it non-obvious. TensorBoard also visualizes model graphs and histograms of weights/gradients (add_histogram), which occasionally reveals dead layers at a glance. Keep the CSV habit alongside it: TensorBoard is for watching, files are for the record.

Running sweeps without chaos

Sooner or later you'll grid-search hyperparameters. The disciplined way is a loop over configs, each producing its own run directory — never hand-edited runs:

import itertools, copy

grid = {"lr": [1e-2, 1e-3], "hidden_dim": [32, 64], "seed": [0, 1]}
keys = list(grid)
for values in itertools.product(*(grid[k] for k in keys)):
    cfg_run = copy.deepcopy(base_cfg)
    for k, v in zip(keys, values): cfg_run["train"][k] = v  # (adjust nesting as needed)
    run_dir = f"runs/sweep_{'_'.join(f'{k}{v}' for k, v in zip(keys, values))}"
    train_with_config(cfg_run, run_dir)   # your Chapter 7 loop + Chapter 9 checkpoints

Two sweep rules from hard experience: (1) vary one thing at a time when diagnosing, grid when exploring — a full grid over 4 hyperparameters is 81 runs and teaches you less than you think; (2) include seeds in the sweep for any comparison that matters (Silent bug C, Chapter 10). After the sweep, the summary table practically writes itself:

lr	hidden_dim	seed	best val acc
1e-2	32	0	0.961
1e-2	32	1	0.958
...	...	...	...

That table — with seeds, from files, reproducible — is what "we performed a hyperparameter sweep" means in a paper. Anything less is hand-waving.

A minimal requirements.txt

Record the environment alongside the code. This is enough to start:

torch>=2.0
torchvision>=0.15
numpy
matplotlib
pyyaml
tensorboard

Generate the exact pinned versions with pip freeze > requirements.txt at publication time. Pinning matters: "torch>=2.0" is a wish; "torch==2.3.1" is a record.

The reproduction README: writing for a stranger

Every research project should have a README.md that lets a stranger reproduce your main result. Not your lab mate — a stranger, with a fresh machine. The template:

# Project title — one-line description

## Results
| Experiment | Val accuracy | Config |
|---|---|---|
| baseline | 0.962 | configs/baseline.yaml |

## Reproduce
1. pip install -r requirements.txt
2. python src/download_data.py        # or: place data in data/
3. python src/train.py --config configs/baseline.yaml
4. python src/evaluate.py --checkpoint runs/<timestamp>/checkpoints/best_model.pt --config configs/baseline.yaml

Expected: ~0.96 val accuracy in ~20 min on a T4 GPU (Colab).

Write it while the project is fresh, and test it by following it yourself on a clean environment (or asking a colleague to). READMEs written from memory months later are fiction; READMEs tested by execution are documentation. Reviewers who can reproduce your work in four commands trust your work — and cite it.

Code review checklist for research code

Before any run whose numbers might be published, review the training script against this list (or better, have a colleague do it — fresh eyes catch what yours skip):

• [ ] Seeds set for random, numpy, torch, CUDA, and data splits — and logged
• [ ] No hard-coded hyperparameters; everything from the config
• [ ] Train/val/test splits are disjoint (no leakage) and seeded
• [ ] Preprocessing identical (modulo augmentation) across splits
• [ ] Loss matches the task; targets have the right shape/dtype
• [ ] model.train()/model.eval() + torch.no_grad() in the right places
• [ ] Checkpoints save (last + best) and the round-trip was tested
• [ ] Metrics logged per epoch; final numbers written to a summary file
• [ ] One-batch overfit test passed on this exact code + data

Nine boxes, five minutes, and the large majority of "embarrassing rebuttal" situations never happen.

Key Takeaways

1. An experiment = code + config + seed + data + environment, all recorded. Anything less is "trying something."
2. Seed everything (random, numpy, torch, CUDA, data splits) and log the seed.
3. Hyperparameters live in config files, never hard-coded — the run command plus config fully describes the experiment.
4. Log per-epoch metrics to CSV and final numbers to a JSON summary; log what goes in the paper's table, when you compute it.
5. Use the src/ + configs/ + runs/ layout; version-control code and configs with git, tag reported results.
6. Run the pre-launch checklist before any paper-bound experiment.

---

Chapter 12: What to Learn Next — A Roadmap

You now speak PyTorch

Take stock: you can create and manipulate tensors, differentiate through arbitrary computations, build networks with nn.Module, choose losses and optimizers, write full training loops with validation, feed real data through custom datasets, checkpoint experiments, debug the ten classic failures, and organize runs reproducibly. That is the complete foundation — everything else in the PyTorch world is a power tool built on top of it, and power tools make sense only once the foundation is solid.

This final chapter maps where to go next, in the order that pays off for a researcher. For each direction: what it is, when you need it, and the honest cost of adopting it.

1. PyTorch Lightning — structure without boilerplate

What: A lightweight framework that organizes your PyTorch code into a LightningModule (model + training step + validation step) and a Trainer (the loop, checkpointing, logging, multi-GPU). Your Chapter 7 loop becomes:

# sketch — the ideas, not a full tutorial
class LitClassifier(pl.LightningModule):
    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = self.loss_fn(self(x), y)
        self.log("train_loss", loss)
        return loss
    def configure_optimizers(self):
        return torch.optim.AdamW(self.parameters(), lr=1e-3)

trainer = pl.Trainer(max_epochs=25, accelerator="gpu", devices=1)
trainer.fit(model, train_loader, val_loader)

When: As soon as your training scripts get repetitive — usually your second or third project. Cost: You must learn its conventions (hooks like training_step); debugging goes through an extra layer. Verdict: worth it for multi-experiment research; skip it while you're still internalizing the raw loop.

2. Hugging Face Transformers (and Diffusers, Datasets)

What: The largest public library of pre-trained models — BERT, GPT-style models, Vision Transformers, Stable Diffusion — with a unified API, plus datasets for one-line dataset loading.

from transformers import AutoModelForSequenceClassification, AutoTokenizer
tok = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=3)

When: The moment your research touches language, vision-language, or generative models — which, statistically, it will. Fine-tuning a pre-trained model beats training from scratch in most student-scale projects. Cost: These are big models with big APIs; start with a fine-tuning tutorial on a small dataset before adapting to your problem. This is also where the tiny learning rates from Chapter 6 (1e-5) become daily practice.

3. torch.compile — free speed

What: PyTorch 2's built-in compiler: one line that fuses and optimizes your model's operations, often 20–50% faster with no code changes:

model = torch.compile(model)  # that's the whole API

When: When training is too slow and you've already fixed data loading (Chapter 8). Cost: First-epoch compilation takes time; rare incompatibilities with dynamic control flow. Try it early — it's nearly free performance.

4. Distributed training — multiple GPUs

What: Training across several GPUs/machines with DistributedDataParallel (DDP): each GPU processes different batches, gradients are averaged. Lightning and Hugging Face's accelerate wrap the considerable setup boilerplate.

When: When one GPU can't finish your experiment in reasonable time — typically large transformers or ImageNet-scale vision. Cost: Real engineering complexity (process groups, samplers, synchronization bugs). Don't touch this until single-GPU training is second nature; most MS theses never need it.

5. Mixed precision, quantization, deployment

• Mixed precision (torch.cuda.amp): trains with 16-bit floats where safe — roughly 2x speed and half the memory on modern GPUs, with minimal accuracy impact. The standard for large-model training.
• Quantization: shrinking a trained model to 8-bit (or less) for inference — how models ship to phones and edge devices.
• Deployment: TorchScript / ONNX export for serving models outside Python.

When: Mixed precision when memory/speed binds at scale; quantization and deployment when your research needs to ship, not just publish. For most early researchers these are "know they exist" topics until a project demands them.

6. JAX and TensorFlow — the informed second framework

From Chapter 1: learn JAX when you want maximum mathematical control and speed (its grad/jit/vmap transformations are genuinely beautiful for methods research), and TensorFlow/Keras when a job or deployment target requires it. Learning either is dramatically easier after PyTorch, because the concepts — tensors, graphs, losses, optimizers — transfer completely; only the API differs.

First steps for each direction

Concretely, here's the smallest useful action for each item above — do these in order, and never all at once:

1. Lightning: Take your Chapter 7 train_one_epoch/evaluate script and convert it to a LightningModule following the official "Lightning in 15 minutes" tutorial. If the conversion feels like it's hiding things you want to see, you're not ready — come back after two more raw-PyTorch projects.
2. Hugging Face: Run the official course's first fine-tuning example (text classification with a small model like DistilBERT) end-to-end on Colab. Then swap in your dataset. The pattern — tokenizer, Trainer, tiny LR — is the same everywhere in this ecosystem.
3. torch.compile: Add the one line to your fastest existing script and measure wall-clock time before/after. If you hit a compilation error, wrap the failing region with torch._dynamo.disable or fall back — don't fight the compiler on your first try.
4. Distributed: Don't start here. When single-GPU runs truly can't finish in time, begin with Hugging Face accelerate (it handles the boilerplate) on 2 GPUs, and read the DDP docs only when you need to understand a failure.
5. Mixed precision: Wrap your optimizer step in torch.cuda.amp following the official AMP recipe on an existing working script. Verify accuracy matches the full-precision run before trusting it — silent degradation is the failure mode.
6. JAX/TensorFlow: After this book, reimplement Chapter 5's classifier in the new framework. You'll be shocked how fast it goes — that's the concepts transferring. Then read one paper's code in that framework to calibrate.

The meta-rule: one new tool per project. Researchers who adopt three new tools at once can't tell which one broke; researchers who add one at a time build a reliable toolkit.

Your 90-day roadmap

Weeks	Focus	Milestone
1–2	Chapters 2–7 of this book	Train an MLP/CNN on MNIST or CIFAR-10 from scratch; hit the expected accuracy
3–4	Chapters 8–11	Reorganize into src/ + configs + runs; reproduce a tiny published result (e.g., a paper's MNIST baseline) and log everything
5–6	torchvision pre-trained models	Fine-tune a pre-trained ResNet on a small dataset of your choice
7–8	Hugging Face	Fine-tune a small transformer (e.g., DistilBERT) on a text classification task
9–12	Your research problem	Apply the full pipeline — data, model, experiment harness — to your thesis/curiosity problem; write it up

The pattern to internalize: foundation → pre-trained models → your problem. Each stage reuses everything before it.

How to keep learning (the researcher way)

• Read paper code, not just papers. For every paper you cite, skim its repository's train.py and model.py. You'll start recognizing this book's patterns everywhere — that's the point.
• Reproduce before you invent. Pick a small published result and reproduce its numbers. It's the fastest way to find gaps in your understanding, and successful reproductions occasionally become workshop papers themselves.
• Teach one concept. Explain autograd (or broadcasting, or the training loop) to a junior student. The questions they ask will reveal what you only half-understood.
• Follow the source. PyTorch's documentation and tutorials (pytorch.org/tutorials) are excellent and current — prefer them over random blog posts, which decay as APIs change.

For Your Research

Your next paper starts with a reproduced baseline. The single most valuable habit for an early researcher: before proposing anything new, reproduce the strongest existing baseline on your problem, with this book's full harness (seeded, logged, checkpointed). A reproduced baseline gives you three things: (1) certainty your pipeline works, (2) a fair comparison point reviewers trust, and (3) deep familiarity with the problem that often suggests the novel idea itself. Most "novel" contributions are born while staring at a baseline's failure cases. The researcher who can train, debug, and reproduce is the researcher whose new ideas actually get tested — and tested ideas become papers.

Key Takeaways

1. The foundation (this book) comes first; everything else — Lightning, Hugging Face, distributed — is a power tool on top of it.
2. Lightning removes boilerplate when scripts get repetitive; Hugging Face is essential for language/vision-language work; torch.compile is nearly free speed.
3. Distributed training, mixed precision, and quantization matter at scale — learn them when a project demands them, not before.
4. JAX/TensorFlow are easier after PyTorch — concepts transfer, only APIs differ.
5. The researcher growth loop: reproduce a baseline → understand its failures → test your idea → write it up.

---

Learning Dashboard

Key APIs at a glance

Task	API	Notes
Create tensor	torch.tensor, torch.randn, torch.zeros/ones, torch.arange	randn = standard normal
Inspect	x.shape, x.dtype, x.device, x.numel()	Print these when debugging
Reshape	x.reshape(...), x.view(...), x.unsqueeze/squeeze, x.T	cat glues existing dim; stack adds new dim
Math	+ - * / @, torch.matmul, .sum/.mean(dim=), .argmax(dim=)	@ = matmul, * = element-wise
Device	x.to(device), x.cpu(), x.cuda()	Model and data must match
NumPy bridge	torch.from_numpy(a), x.numpy()	Shared memory; CPU only
Track gradients	requires_grad=True, x.backward(), x.grad	Backward needs scalar
Pause tracking	torch.no_grad(), x.detach()	Eval, logging, manual updates
Build model	nn.Module, nn.Linear, nn.Sequential, nn.Parameter	Call model(x), not forward
Activations	nn.ReLU, nn.Sigmoid, nn.Tanh, nn.Softmax(dim=1)	ReLU default for hidden
Losses	nn.MSELoss, nn.BCEWithLogitsLoss, nn.CrossEntropyLoss	Logits in, no softmax first
Optimizers	torch.optim.SGD, Adam, AdamW	zero_grad → backward → step
Schedulers	StepLR, CosineAnnealingLR, ReduceLROnPlateau	Decay LR during training
Data	Dataset, DataLoader, TensorDataset, random_split	Shuffle train, seed splits
Save/load	torch.save/load, model.state_dict(), load_state_dict	Save dicts, not objects
Reproducibility	torch.manual_seed, random.seed, np.random.seed	Log the seed

Tensor shape conventions

Data	Shape	Example
Scalar (loss)	()	loss.item() → Python float
One sample, F features	(F,)	(8,)
Batch of samples	(B, F)	(64, 8)
Batch of images	(B, C, H, W)	(64, 3, 32, 32)
Batch of sequences	(B, T, F)	(32, 50, 128)
Linear layer weights	(out, in)	nn.Linear(10, 32) → (32, 10)
Class logits	(B, K)	(64, 3) for 3 classes
Class targets	(B,) int64	values 0..K-1

Which loss / optimizer when

Task	Loss	Output	Targets
Regression	MSELoss	(N,1) floats	(N,1) floats
Binary classification	BCEWithLogitsLoss	(N,1) logits	(N,1) float 0/1
Multi-class (K)	CrossEntropyLoss	(N,K) logits	(N,) int64 indices
Multi-label	BCEWithLogitsLoss	(N,K) logits	(N,K) float 0/1

Situation	Optimizer	LR starting point
Default / prototype / transformer	AdamW	1e-3 (1e-4 for transformers)
CNN from scratch	SGD + momentum 0.9	0.1 with scheduler
Fine-tuning pre-trained	AdamW	1e-5 – 5e-5

Hyperparameter starting points

Hyperparameter	Start	If loss explodes/NaN	If loss barely moves
Learning rate	1e-3 (AdamW)	÷ 10	× 3–10
Batch size	64	÷ 2 (also fixes OOM)	× 2 (faster/epoch)
Weight decay	1e-4	—	—
Epochs	until val loss plateaus	—	train longer / bigger model

Debugging quick reference

Symptom	Check first
Shape error	Print shapes at each stage; read the message literally
Device mismatch	One device var; .to(device) on model + every batch
Weird loss, no crash	Missing zero_grad(); sigmoid/softmax before loss
NaN loss	LR too high (÷10); check for log(0)/div-by-zero
CUDA OOM	Halve batch size; gradient accumulation
Noisy test accuracy	Missing model.eval() + torch.no_grad()
Slow training	num_workers, pin_memory, profile samples/sec
"Can it learn at all?"	One-batch overfit test

Shape-debugging decision tree

When a shape error strikes, walk this tree:

Error mentions shapes (AxB) vs (CxD)?
├─ In a Linear layer → does input's last dim == layer's in_features?
│   └─ No → fix the layer definition or the previous layer's out_features
├─ In a loss function → do output and target shapes match the loss's contract?
│   └─ Check the "Which loss / optimizer when" table above (logits vs indices!)
├─ In torch.cat/stack → are all non-concat dims equal? (cat) / all dims equal? (stack)
│   └─ No → you probably want the other one, or a reshape first
├─ After a transpose/permute → is a later .view() failing?
│   └─ Yes → use .reshape() or .contiguous().view()
└─ Shapes "almost" match (e.g. (32,1) vs (32,)) → .squeeze()/.unsqueeze() to align

Print x.shape after every transformation while learning. Professionals do this too — they just delete the prints afterward.

Complete minimal training template (copy-paste starter)

This is Chapters 2–11 compressed into one file. Start every new project from it:

"""Minimal PyTorch training template — AstolixGen Learning Series, Book 11."""
import argparse, csv, json, os, random, time
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset, random_split

# ---------- setup ----------
def set_seed(seed=42):
    random.seed(seed); np.random.seed(seed)
    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)

def get_args():
    p = argparse.ArgumentParser()
    p.add_argument("--lr", type=float, default=1e-2)
    p.add_argument("--batch-size", type=int, default=64)
    p.add_argument("--epochs", type=int, default=25)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--run-dir", default=None)
    return p.parse_args()

# ---------- model ----------
class MLP(nn.Module):
    def __init__(self, in_dim=8, hidden=64, out_dim=3):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(in_dim, hidden), nn.ReLU(),
                                 nn.Linear(hidden, hidden), nn.ReLU(),
                                 nn.Linear(hidden, out_dim))
    def forward(self, x):
        return self.net(x)

# ---------- loop ----------
def run_epoch(model, loader, loss_fn, device, optimizer=None):
    train = optimizer is not None
    model.train(train)
    tot_loss, correct, total = 0.0, 0, 0
    ctx = torch.enable_grad() if train else torch.no_grad()
    with ctx:
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            if train: optimizer.zero_grad()
            out = model(xb)
            loss = loss_fn(out, yb)
            if train:
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
            tot_loss += loss.item() * xb.size(0)
            correct += (out.argmax(1) == yb).sum().item()
            total += xb.size(0)
    return tot_loss / total, correct / total

def main():
    args = get_args(); set_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    run_dir = args.run_dir or f"runs/{time.strftime('%Y%m%d_%H%M%S')}"
    os.makedirs(f"{run_dir}/checkpoints", exist_ok=True)
    json.dump(vars(args), open(f"{run_dir}/config.json", "w"), indent=2)

    torch.manual_seed(args.seed)
    X = torch.randn(3000, 8); y = (X[:, :3].argmax(1))
    tr, va = random_split(TensorDataset(X, y), [2400, 600],
                          generator=torch.Generator().manual_seed(args.seed))
    tr_loader = DataLoader(tr, batch_size=args.batch_size, shuffle=True)
    va_loader = DataLoader(va, batch_size=256)

    model = MLP().to(device)
    loss_fn = nn.CrossEntropyLoss()
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)

    best, wait = float("inf"), 0
    with open(f"{run_dir}/metrics.csv", "w", newline="") as f:
        w = csv.writer(f); w.writerow(["epoch", "train_loss", "train_acc", "val_loss", "val_acc"])
        for epoch in range(1, args.epochs + 1):
            tr_l, tr_a = run_epoch(model, tr_loader, loss_fn, device, opt)
            va_l, va_a = run_epoch(model, va_loader, loss_fn, device)
            w.writerow([epoch, f"{tr_l:.4f}", f"{tr_a:.4f}", f"{va_l:.4f}", f"{va_a:.4f}"])
            print(f"ep {epoch:3d} train {tr_l:.4f}/{tr_a:.3f} val {va_l:.4f}/{va_a:.3f}")
            torch.save({"epoch": epoch, "model": model.state_dict(),
                        "opt": opt.state_dict(), "best": best},
                       f"{run_dir}/checkpoints/last.pt")
            if va_l < best:
                best, wait = va_l, 0
                torch.save(model.state_dict(), f"{run_dir}/checkpoints/best_model.pt")
            else:
                wait += 1
                if wait >= 10:
                    print("early stopping"); break
    json.dump({"best_val_loss": best, "seed": args.seed,
               "torch": torch.__version__}, open(f"{run_dir}/summary.json", "w"), indent=2)
    print("done ->", run_dir)

if __name__ == "__main__":
    main()

Run it: python template.py --lr 1e-2 --epochs 25. It includes seeds, configs-as-flags, CSV logging, rolling + best checkpoints, early stopping, gradient clipping, and device handling — the whole book in 100 lines.

Paper-to-code worksheet

When reading a paper with code, fill this in. If you can't fill a row, that's what to ask the authors (or what to figure out next):

Paper says	Code has	This book's chapter
Architecture diagram	nn.Module class in model.py	Ch. 5
"Trained with cross-entropy"	loss_fn = nn.CrossEntropyLoss()	Ch. 6
"Adam, lr 3e-4"	torch.optim.AdamW(params, lr=3e-4)	Ch. 6
"Batch size 256, 100 epochs"	DataLoader(..., batch_size=256), epoch loop	Ch. 7/8
Dataset description	Dataset class + transforms	Ch. 8
"Best checkpoint selected by validation"	best_model.pt logic	Ch. 9
Reported numbers	summary.json / results table	Ch. 11
"Code available"	Tagged git commit + README	Ch. 11

NumPy → PyTorch cheat sheet

If you know NumPy, you know most tensor ops — the names transfer almost directly:

NumPy	PyTorch	Notes
np.array([...])	torch.tensor([...])	Same idea
np.zeros((2,3))	torch.zeros(2, 3)	Tuple optional in torch
np.random.randn(2,3)	torch.randn(2, 3)	Same distribution
a.reshape(3, 4)	a.reshape(3, 4)	Identical
a.T	a.T (2-D) / a.permute(...) (N-D)	permute generalizes transpose
a + b (broadcast)	a + b (broadcast)	Same rules
a @ b	a @ b / torch.matmul(a, b)	Same
a * b	a * b	Element-wise in both
np.mean(a, axis=0)	a.mean(dim=0)	axis → dim
np.concatenate([a,b], axis=0)	torch.cat([a, b], dim=0)	Same idea
np.stack([a,b])	torch.stack([a, b])	Same idea
a[1:3, :]	a[1:3, :]	Identical slicing
a[a > 0]	a[a > 0]	Identical boolean masking
np.argmax(a, axis=1)	a.argmax(dim=1)	Same

The mental model: PyTorch tensors are NumPy arrays with .to(device) and autograd. When stuck, ask "how would I do this in NumPy?" — the answer usually works verbatim.

---

Appendix A: Error Message Catalog

A quick-scan table of the exact error strings you'll meet, what they mean, and the one-line fix. (Chapter 10 explains each in depth.)

Error message (key phrase)	Meaning	Fix
mat1 and mat2 shapes cannot be multiplied (AxB and CxD)	Linear layer got wrong input features	Match layer in_features to data's last dim
Expected all tensors to be on the same device	Model/data split across CPU/GPU	One device var; .to(device) everywhere
expected scalar type Long but found Float (or reverse)	Target dtype wrong for the loss	.long() for CE targets; .float() for BCE/MSE
Expected target size (N,), got (N, 1)	Target has an extra dim	.squeeze() the target
grad can be implicitly created only for scalar outputs	.backward() on non-scalar	Reduce to scalar first (your loss does this)
one of the variables needed for gradient computation has been modified by an inplace operation	In-place op corrupted the graph	Use out-of-place ops; avoid inplace=True while debugging
CUDA out of memory	Exceeded GPU RAM	Halve batch size; gradient accumulation
view size is not compatible with input tensor's size and stride	view after transpose/permute	Use .reshape() or .contiguous().view()
DataLoader worker exited unexpectedly	Multiprocessing failure	Debug with num_workers=0 first
stack expects each tensor to be equal size	Variable-size samples in default collate	Write a custom collate_fn (Ch. 8)
Missing key(s) in state_dict / Unexpected key(s)	Architecture/checkpoint mismatch	Recreate exact architecture; or strict=False deliberately
Attempting to deserialize object on a CUDA device	GPU checkpoint on CPU-only machine	torch.load(path, map_location="cpu")
Loss is nan (no exception)	Divergence or bad data	÷10 LR; check data for NaNs; NaN-detect in loop

Appendix B: Further Reading (Real Resources)

These are well-known, maintained resources worth your time after this book — all real, all free or widely available:

• The official PyTorch tutorials (pytorch.org/tutorials) — the single best next step; the "60-minute blitz" and topic guides are kept current with each release.
• Stevens, Antiga & Viehmann, Deep Learning with PyTorch (Manning) — the canonical book-length treatment; goes deeper on everything here plus deployment.
• The Hugging Face NLP/course materials (huggingface.co/learn) — free, hands-on courses for transformers and diffusion, built on PyTorch.
• fast.ai's Practical Deep Learning (course.fast.ai) — a top-down, code-first complement to this book's bottom-up approach.
• Andrej Karpathy's "Zero to Hero" series (public video lectures + code) — builds neural nets from scratch to GPT; superb for intuition.
• PyTorch Lightning documentation (lightning.ai/docs) — when you're ready for Chapter 12's first power tool.
• Papers with Code (paperswithcode.com) — find any paper's official PyTorch implementation to practice the paper-to-code mapping from Chapter 1.

Prefer primary sources (official docs, the paper itself, the authors' repo) over random tutorials — the ecosystem moves fast and third-party posts decay.

References

[1] A. Paszke et al., "PyTorch: An imperative style, high-performance deep learning library," in Advances in Neural Information Processing Systems 32 (NeurIPS 2019), Vancouver, Canada, 2019.

[2] D. P. Kingma and J. Ba, "Adam: A method for stochastic optimization," in Proc. 3rd Int. Conf. Learning Representations (ICLR 2015), San Diego, CA, USA, 2015.

[3] D. E. Rumelhart, G. E. Hinton, and R. J. Williams, "Learning representations by back-propagating errors," Nature, vol. 323, no. 6088, pp. 533–536, 1986.

[4] Y. LeCun, Y. Bengio, and G. Hinton, "Deep learning," Nature, vol. 521, no. 7553, pp. 436–444, 2015.

[5] A. Krizhevsky, I. Sutskever, and G. E. Hinton, "ImageNet classification with deep convolutional neural networks," in Advances in Neural Information Processing Systems 25 (NeurIPS 2012), Lake Tahoe, NV, USA, 2012.

[6] K. He, X. Zhang, S. Ren, and J. Sun, "Deep residual learning for image recognition," in Proc. IEEE Conf. Computer Vision and Pattern Recognition (CVPR 2016), Las Vegas, NV, USA, 2016, pp. 770–778.

[7] S. Ioffe and C. Szegedy, "Batch normalization: Accelerating deep network training by reducing internal covariate shift," in Proc. 32nd Int. Conf. Machine Learning (ICML 2015), Lille, France, 2015, pp. 448–456.

[8] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov, "Dropout: A simple way to prevent neural networks from overfitting," J. Machine Learning Research, vol. 15, no. 1, pp. 1929–1958, 2014.

[9] X. Glorot and Y. Bengio, "Understanding the difficulty of training deep feedforward neural networks," in Proc. 13th Int. Conf. Artificial Intelligence and Statistics (AISTATS 2010), Sardinia, Italy, 2010, pp. 249–256.

[10] I. Loshchilov and F. Hutter, "Decoupled weight decay regularization," in Proc. 7th Int. Conf. Learning Representations (ICLR 2019), New Orleans, LA, USA, 2019.

[11] H. Robbins and S. Monro, "A stochastic approximation method," Annals of Mathematical Statistics, vol. 22, no. 3, pp. 400–407, 1951.

[12] A. Vaswani et al., "Attention is all you need," in Advances in Neural Information Processing Systems 30 (NeurIPS 2017), Long Beach, CA, USA, 2017.

---

Glossary

• Activation function — A non-linear function (e.g., ReLU) applied after a layer so the network can learn non-linear patterns.
• Autograd — PyTorch's automatic differentiation engine; records operations and computes gradients via the chain rule.
• Backward pass — The phase of training where gradients are computed by traversing the computation graph in reverse (.backward()).
• Batch — A group of samples processed together in one forward/backward pass.
• Batch size — The number of samples per batch; a key hyperparameter affecting speed, memory, and generalization.
• Broadcasting — PyTorch's rule for operating on differently-shaped tensors by stretching size-1 dimensions.
• Checkpoint — A saved file containing model weights plus optimizer/scheduler state and epoch, allowing exact training resumption.
• Computation graph — The recorded chain of operations built during the forward pass, used by autograd for the backward pass.
• Cross-entropy loss — The standard loss for multi-class classification; rewards high predicted probability for the correct class.
• CUDA — NVIDIA's platform for running computations on GPUs; PyTorch uses it to accelerate tensor operations.
• DataLoader — PyTorch utility that wraps a Dataset to yield shuffled, batched, optionally parallel-loaded samples.
• Dataset — A class defining __len__ and __getitem__ that provides access to individual training samples.
• Device — Where a tensor lives and computes: CPU or a GPU (cuda:0).
• Dropout — A regularization technique that randomly zeroes activations during training to reduce overfitting.
• dtype — A tensor's data type (e.g., float32, int64); must be consistent across operations.
• Epoch — One complete pass through the training dataset.
• Forward pass — Running input data through the model to produce predictions.
• Gradient — How much the loss would change if a parameter changed slightly; the signal used to update weights.
• Gradient accumulation — Summing gradients over several small batches before stepping, simulating a larger batch size.
• Learning rate — The step size of the optimizer; the most influential hyperparameter.
• Loss function — A function producing a single number measuring how wrong the model's predictions are.
• Mini-batch — A small batch (see batch); stochastic gradient descent trains on mini-batches.
• Mixed precision — Training with 16-bit floats where safe, for speed and memory savings on modern GPUs.
• nn.Module — PyTorch's base class for all neural network components; holds layers, parameters, and the forward pass.
• Optimizer — The algorithm (SGD, Adam, AdamW) that converts gradients into weight updates.
• Overfitting — When a model memorizes training data but generalizes poorly to new data.
• Parameter — A learnable tensor inside a module (weights, biases), updated by the optimizer.
• ReLU — Rectified linear unit, max(0, x); the default hidden-layer activation function.
• requires_grad — A tensor flag telling autograd to track operations for gradient computation.
• Scheduler — Adjusts the learning rate during training (e.g., decay on a schedule).
• Seed — A fixed starting point for random number generators, enabling reproducible runs.
• Softmax — A function converting raw scores into probabilities that sum to 1.
• state_dict — A dictionary mapping layer names to weight tensors; the standard format for saving models.
• Tensor — PyTorch's multi-dimensional array; the core data structure for data, weights, and gradients.
• Validation set — Data held out from training, used to honestly measure generalization during development.
• Weight decay — L2 regularization added to the optimizer, gently pulling weights toward zero to fight overfitting.
• Ablation study — An experiment removing or changing one component at a time to measure its contribution; standard in papers.
• Attention mask — A tensor marking which positions are real data vs. padding, so models ignore padded positions.
• Batch normalization — A layer normalizing activations per batch during training (running statistics at eval); stabilizes deep training.
• Checkpoint (best/last) — best_model.pt: weights with the best validation score (for reporting); last.pt: latest state (for resuming).
• Collate function — DataLoader hook assembling individual samples into a batch; customized for variable-length data.
• Confusion matrix — A table of true vs. predicted classes revealing which mistakes a model makes.
• Determinism — Property of producing identical results across runs given the same seed; partial on GPUs.
• Early stopping — Halting training when validation stops improving, keeping the best weights; simple anti-overfitting.
• Fine-tuning — Continuing training of a pre-trained model on a new task, usually with a small learning rate.
• Generalization — A model's ability to perform well on unseen data (as opposed to memorized training data).
• Gradient clipping — Capping the gradient norm before the optimizer step to prevent explosive updates.
• Inference — Using a trained model to make predictions (no training, no gradients).
• Kaiming/He initialization — Weight initialization scaled for ReLU networks, keeping activation variance stable.
• Label smoothing — Softening hard 0/1 targets slightly to discourage overconfident predictions; a mild regularizer.
• Learning-rate scheduler — (see Scheduler) — e.g., step decay, cosine annealing, plateau reduction.
• Logit — A raw, unnormalized model output score (before softmax/sigmoid); what losses expect as input.
• Model card — Documentation accompanying shared weights: architecture, data, intended use, limitations.
• One-batch overfit test — Debugging check: a working setup should memorize a single batch quickly.
• Positional encoding — Information about sequence position added to inputs (often a buffer); used in transformers.
• Pre-trained model — A model trained on a large dataset, reused as a starting point (feature extractor or fine-tuning base).
• Regularization — Techniques (weight decay, dropout, augmentation) reducing overfitting by constraining the model.
• Teacher forcing / transfer learning — (see Fine-tuning) — adapting pre-trained knowledge to a new task.
• Test set — Data touched once, at the very end, for the final unbiased performance estimate.
• Underfitting — When a model is too simple (or undertrained) to capture even the training data's patterns.
• Xavier/Glorot initialization — Weight initialization keeping signal variance stable across layers; suited to tanh/sigmoid.

---

Practice Exercises

1. Environment check. On Colab (or your machine), print the PyTorch version, whether CUDA is available, and the GPU name. Write one paragraph explaining in your own words why the same code runs on CPU and GPU with only a device change.
2. Tensor gymnastics. Create a random tensor of shape (16, 10). Without Python loops: compute per-column means, subtract them, then verify the result has (near-)zero column means. Print shapes at each step.
3. Broadcasting detective. Predict the output shapes of: (4, 1) + (3,), (5, 4, 1) * (4, 7), (2, 3) + (2, 1, 3). Then verify with code. For any that surprise you, write down the broadcasting rule that explains them.
4. Autograd by hand. For y = (x1 * x2 + x3)^2 with x = [2.0, 3.0, 1.0], compute the gradients analytically on paper, then verify with PyTorch autograd. They must match.
5. Build and count. Implement a 3-layer MLP (20 → 64 → 64 → 5) as an nn.Module subclass. Count its parameters with the one-liner from Chapter 5 and verify the count by hand calculation.
6. Loss matching (research-oriented). You are given a dataset with 4 possible tags per sample, where each sample can have multiple tags. Which loss function do you use, what shape/dtype are the model outputs and targets, and why is CrossEntropyLoss wrong here? Implement a 50-line training script proving your choice works on synthetic data.
7. The one-batch test (research-oriented). Take the Chapter 7 training script and modify it to train on a single batch of 32 samples. Show that the loss goes to ~0. Then deliberately introduce one bug from Chapter 10 (e.g., remove zero_grad()) and document how the symptom differs. Write up both runs as a one-page lab note.
8. Custom dataset (research-oriented). Build a Dataset class for a small real dataset of your choice (a CSV you create, or MNIST via torchvision). Include a seeded train/val split, correct dtypes in __getitem__, and train the Chapter 7 classifier on it. Log per-epoch metrics to a CSV file.
9. Reproduce and log (research-oriented). Reproduce a tiny published-style experiment: train an MLP on MNIST to ≥ 97% test accuracy using the full harness from Chapters 7–11 (seeded, config file, metrics CSV, best_model.pt, git commit). Write the summary.json you would paste into a paper's results table, including PyTorch version and seed.
10. Ablation study (research-oriented). Using your Exercise 9 setup, run three ablations changing exactly one thing each: (a) SGD with momentum instead of AdamW, (b) no hidden layer (linear model), (c) learning rate 10× higher. Record all results in a table with seeds and configs. Write three sentences interpreting the table — this is the core skill behind every "ablation study" section in published papers.

---

End of Book 11 — Deep Learning with PyTorch: First Steps AstolixGen Learning Series · 50 books · written for researcher students