The Deep Learning Mental Model
PyTorch's tensors, its autograd engine, and the training loop can look like three separate topics learned one after another.
They are actually one machine: a computation graph built from tensor operations, a backward pass that walks that graph to compute gradients, and a training loop that repeatedly uses those gradients to nudge parameters toward lower loss.
This page connects PyTorch Basics, Autograd, Training Loops, and GPUs & Mixed Precision into one model, rather than four separate topics to memorize independently.
Summary
- Every PyTorch training run is the same loop - build a computation graph via a forward pass, walk it backward to get gradients, and use those gradients to update parameters - and everything else in this section is a variation or optimization of that one loop.
- Insight: Understanding the loop as one mechanism, rather than a sequence of API calls to memorize, is what lets you actually debug a model that isn't learning instead of guessing at fixes.
- Key Concepts: tensor, computation graph, autograd, backward pass, optimizer step, mixed precision.
- When to Use: Designing or debugging any custom training loop, deciding whether GPU or mixed precision is worth the added complexity, or reasoning about why a model's loss isn't decreasing.
- Limitations/Trade-offs: The computation graph and stored activations consume memory roughly proportional to model depth and batch size, and mixed precision trades numerical precision for speed and memory savings.
- Related Topics: transfer learning, distributed training, PyTorch Lightning.
Foundations
A tensor is PyTorch's core data structure, an n-dimensional array much like a NumPy array, but with two additions that matter for deep learning: it can live on a GPU, and it can track the operations applied to it.
That second property is what makes autograd possible: every operation performed on a tensor with requires_grad=True is recorded into a computation graph, a directed graph where nodes are operations and edges are the tensors flowing between them.
A useful analogy is a recipe with a receipt attached.
The forward pass is following the recipe, mixing tensors through operations to produce an output, and the receipt is the computation graph, listing exactly which operations happened in which order so you can later walk backward through it to see how much each ingredient contributed to the final result.
The forward pass is ordinary computation: input tensors flow through layers such as matrix multiplies, activations, and normalization, to produce an output like a predicted class or a next-token probability.
The backward pass is where autograd's real function shows up.
Calling .backward() on a scalar loss walks the recorded graph in reverse, applying the chain rule at each node to compute the gradient of the loss with respect to every parameter that contributed to it.
Mechanics & Interactions
The training loop is this pair of passes repeated many times, and every step in it exists to answer a specific question in a fixed order:
for each batch:
forward pass -> what did the model predict? (build graph)
compute loss -> how wrong was the prediction?
backward pass -> how much did each parameter contribute to that error? (walk graph)
optimizer step -> nudge each parameter opposite its gradient
zero gradients -> clear old gradients before the next batch
The optimizer step - SGD, Adam, and their variants - is where the actual learning happens, taking the gradients autograd computed and updating each parameter by some function of that gradient, typically scaled by a learning rate.
Zeroing gradients exists because PyTorch accumulates gradients into .grad by default rather than overwriting them.
That design supports techniques like gradient accumulation across multiple small batches, but it silently corrupts a normal training loop if you forget to clear gradients before the next backward call.
for batch in loader:
optimizer.zero_grad() # clear last step's gradients
output = model(batch.x) # forward pass builds the graph
loss = loss_fn(output, batch.y)
loss.backward() # walk the graph backward
optimizer.step() # update parameters using .gradThe computation graph is rebuilt from scratch on every forward pass and discarded after .backward() by default, which is why PyTorch is described as define-by-run.
The graph is a record of what actually happened on this specific forward pass, not a static structure declared up front, which is what makes data-dependent control flow straightforward compared to frameworks requiring a graph defined before any data flows through it.
GPUs matter here because the forward and backward passes are almost entirely matrix multiplications and element-wise operations, which parallelize extremely well across a GPU's thousands of cores compared to a CPU's much smaller number of cores optimized for sequential work.
Mixed precision training exploits the fact that most of this arithmetic doesn't need full 32-bit floating point accuracy to converge.
Running matrix multiplies in 16-bit, or bfloat16, while keeping a 32-bit master copy of the parameters cuts memory use and often roughly doubles throughput on modern GPU tensor cores, with a gradient scaler used to prevent small gradient values from underflowing to zero in the lower-precision format.
Advanced Considerations & Applications
At scale, the computation graph's memory cost becomes the binding constraint rather than raw compute, since every intermediate activation from the forward pass has to stay in memory until the backward pass consumes it.
Deeper models and larger batch sizes cost memory roughly linearly in both dimensions as a result.
Gradient checkpointing trades compute for memory directly: instead of storing every intermediate activation, it stores only a subset and recomputes the rest during the backward pass, a deliberate application of the same forward/backward mental model rather than a separate technique.
Transfer learning and fine-tuning reuse this model too.
Freezing a pretrained network's early layers with requires_grad=False removes those parameters from the graph's backward walk entirely, so gradients are only computed - and only optimizer steps applied - for the layers actually being adapted, which is what makes fine-tuning cheaper than training from scratch.
Distributed training extends the same loop across multiple GPUs or machines: each device runs its own forward and backward pass on a slice of the batch, and the gradients, not the data or the model, are what gets synchronized between devices before each optimizer step.
This is why network bandwidth for gradient synchronization, not raw compute, is often the bottleneck at scale.
| Precision / Scale Strategy | Strength | Weakness | Best Fit |
|---|---|---|---|
| Full 32-bit precision, single GPU | Simplest to reason about; no numerical surprises | Slowest option; caps model/batch size to one GPU's memory | Small models, debugging, establishing a correctness baseline |
| Mixed precision, single GPU | Meaningful speedup and memory savings on modern tensor cores | Adds a gradient scaler and occasional numerical instability to reason about | Most production training once correctness is established |
| Distributed (multi-GPU) training | Scales model and batch size beyond one device's memory | Gradient synchronization becomes a real bottleneck; adds operational complexity | Models or datasets too large for one GPU regardless of precision |
Common Misconceptions
- "The computation graph is built once when the model is defined." It's rebuilt on every forward pass because PyTorch is define-by-run, not once at model-definition time - this is what lets a model's control flow depend on its input.
- "Calling
.backward()twice in a row just recomputes the same gradients." Without clearing them first, gradients accumulate into.gradrather than being overwritten, so a missingzero_grad()silently corrupts a standard training loop. - "Mixed precision just makes training faster with no downside." It trades some numerical precision for speed and memory, which is why a gradient scaler is needed to prevent small gradients from underflowing to zero in 16-bit formats.
- "More GPUs always means proportionally faster training." Distributed training adds gradient-synchronization overhead between devices, so returns diminish, and can even reverse, once that synchronization cost rivals the compute time it's meant to save.
- "Freezing layers during fine-tuning is only about saving time." It also removes those layers from the backward pass's gradient computation entirely, which is what makes fine-tuning both faster and less prone to overwriting pretrained weights unintentionally.
FAQs
What actually happens when I call loss.backward()?
PyTorch walks the computation graph recorded during the forward pass, in reverse, applying the chain rule at each operation to compute the gradient of the loss with respect to every tensor that has requires_grad=True and contributed to that loss.
Why do I need to call optimizer.zero_grad() every batch?
PyTorch accumulates new gradients into a tensor's existing .grad attribute by default rather than overwriting it. Without clearing it first, gradients from the previous batch would still be added into the current batch's gradients, corrupting the update.
What is a computation graph, concretely?
A directed graph recorded during the forward pass, where each node is an operation applied to a tensor and each edge is a tensor flowing between operations. It's what autograd walks backward to compute gradients, and it's rebuilt from scratch on every forward pass.
Why is PyTorch described as "define-by-run"?
Because the computation graph is built dynamically as operations actually execute on real data, rather than declared statically before any data flows through it. This makes data-dependent control flow, like a different number of loop iterations per input, straightforward to express.
What does the optimizer actually do with the gradients autograd computes?
It updates each parameter using some function of its gradient and a learning rate. SGD subtracts a scaled gradient directly, while adaptive optimizers like Adam also track per-parameter running statistics of past gradients to adjust the effective step size.
Why does GPU training help so much for neural networks specifically?
Forward and backward passes are dominated by matrix multiplications and element-wise operations, which parallelize extremely well across a GPU's large number of simple cores, unlike more sequential workloads that suit a CPU's smaller number of powerful cores better.
What problem does mixed precision training actually solve?
It reduces memory use and increases throughput by running most arithmetic in a lower-precision format, 16-bit or bfloat16, instead of full 32-bit floats, which modern GPU tensor cores execute significantly faster, while keeping a 32-bit master copy of parameters for numerical stability.
Why does mixed precision need a gradient scaler?
Small gradient values can underflow to zero when represented in 16-bit floating point, silently losing information. A gradient scaler multiplies the loss by a scaling factor before the backward pass, and unscales gradients before the optimizer step, to keep values in a representable range.
How does freezing layers during fine-tuning actually save computation?
Setting requires_grad=False on a layer's parameters removes it from the backward pass's gradient computation entirely. Autograd never computes a gradient for it and the optimizer never updates it, so both the backward pass and the update step are cheaper.
What actually gets synchronized in distributed training across multiple GPUs?
Each device computes its own forward and backward pass on its slice of the batch, and it's the resulting gradients, not the raw data or the model weights mid-training, that get synchronized, typically averaged, across devices before each optimizer step.
What is gradient checkpointing and why trade compute for memory?
It stores only a subset of forward-pass activations instead of all of them, recomputing the rest during the backward pass when needed. It's useful when a model is memory-bound rather than compute-bound, since GPU memory is often the harder limit to work around.
Is a bigger batch size always better for training throughput?
Not unconditionally. Larger batches better utilize GPU parallelism up to a point, but they also increase the memory needed to store activations for the backward pass, and very large batches can change training dynamics rather than simply training faster for free.
Related
- PyTorch Basics - tensors, devices, and the core API this model builds on
- Autograd - a deeper look at the backward pass and gradient computation
- Training Loops - the loop this page describes, written out in working code
- GPUs & Mixed Precision - hands-on precision and device placement
- Transfer Learning & Fine-Tuning - freezing layers as an application of this model
- Distributed Training - extending the loop across multiple devices
Stack versions: This page was written for PyTorch 2.6+ and Python 3.14 (stable) / 3.13 (maintenance).