Chapter 6: PyTorch Fundamentals

Every arithmetic operator (+, *, /, torch.exp) works element-wise: output entry $(i, j)$ depends only on input entry $(i, j)$, so a whole batch runs in one kernel. A reduction collapses a dimension instead. In the formula, $X$ is an $n \times d$ matrix (rows are examples, columns are features), $j$ runs over the $d$ columns, and $s_i$ is row $i$'s total, so X.sum(dim=1) has shape $(n,)$. dim names the axis that disappears, exactly like NumPy's axis. Passing keepdim=True leaves it as size 1, giving $(n, 1)$, which keeps later arithmetic aligned. Failure mode: mean() on an integer tensor raises a RuntimeError instead of rounding.

Slicing with integers and ranges, such as L[:, 1:3], returns a view: nothing is copied, so writing into it changes the original. Advanced indexing with a boolean mask (L[L > 1.5]) or an integer list returns a copy, and a mask always yields a flat 1-D result. gather picks one entry per position along dim: in the formula, input is the source tensor, index is an integer tensor with the same number of dimensions, and $\text{out}$ takes index's shape. Failure mode: a (3,) label vector against a 2-D input raises an error until you unsqueeze(1) it. Slicing syntax matches NumPy's.

Without broadcasting, comparing $n$ points to $m$ points needs a double loop or two manually tiled copies. PyTorch follows NumPy's rules exactly: align shapes from the right; each pair of sizes must be equal or one must be 1, and the size-1 side is stretched. In the formula, $n$ and $m$ are the two point counts and $d$ the coordinate dimension; inserting size-1 axes with None or unsqueeze lets every point of one set meet every point of the other. The stretch is virtual: expand sets that stride to 0. Failure mode: compatible-but-wrong shapes never error, so a (3, 1) prediction minus a (3,) target silently becomes (3, 3).

Image loaders often produce channels-last $(B, H, W, C)$ batches while convolutions expect $(B, C, H, W)$, and copying every batch just to reorder axes wastes memory. Recall from the previous topic that a tensor is storage plus strides, the elements skipped per step along each dim. Here $B$ is batch size, $H$ and $W$ height and width, $C$ channels. permute only reorders shape and strides, so the result is a view sharing storage, no longer contiguous. view refuses strides it cannot reinterpret, reshape copies when it must, .contiguous() forces the copy. The risk: an in-place op such as add_ (trailing underscore) writes through every view of that storage.

Linear layers and attention scores are sums of products over a shared dimension, which element-wise * cannot do; @ (torch.matmul) contracts the last dim of the left operand with the second-to-last of the right. In the formula, $N$ is the batch size, each $A[b]$ is $n \times m$, each $B[b]$ is $m \times p$, and $j$ runs over the shared inner size $m$: the row-times-column rule from Linear Algebra. torch.bmm demands exactly 3-D inputs with equal batch; matmul also broadcasts leading dims. torch.einsum spells out indices: in 'bnm,bmp->bnp' a letter missing from the output is summed. Mismatched inner sizes raise mat1 and mat2 shapes cannot be multiplied.

Without a record of how the loss was computed, nothing is differentiable. Setting requires_grad=True on a floating-point tensor tells PyTorch to log every operation that touches it. Each result gets a grad_fn, a node such as MulBackward0 that remembers its inputs and how to differentiate itself. In the formula, $x$ is the input, $y = g(x)$ an intermediate and $z = f(y)$ the output; $\partial z / \partial x$ is the product of two local derivatives, the chain rule from Mathematical Foundations. The graph is rebuilt every forward pass, so Python if statements and loops just work. Integer tensors cannot require gradients; without the flag, .backward() raises an error.

.backward() does not derive a symbolic formula; it runs reverse-mode automatic differentiation. Here $\mathbf{x}$ is an input with $n$ entries, $\mathbf{y}$ an output with $m$ entries, $J$ their $m \times n$ Jacobian, $\bar{\mathbf{y}}$ the upstream gradient shaped like $\mathbf{y}$, and $\bar{\mathbf{x}}$ the result shaped like $\mathbf{x}$. Each grad_fn computes this product for its own operation and the engine chains them back to the leaves: the chain rule as Linear Algebra's matrix products. A scalar loss is seeded with $\bar{\mathbf{y}} = 1$. A non-scalar output has no implied seed, so y.backward() raises grad can be implicitly created only for scalar outputs; pass a gradient shaped like $\mathbf{y}$, or call .sum() first.

Where do gradients land? Only on leaf tensors: ones you created directly with requires_grad=True, like weights, whose grad_fn is None. In the formula, $L$ is the scalar loss, $x$ a leaf and x.grad its stored gradient, which backward adds to rather than overwrites. Intermediate results are non-leaves; their .grad stays None, with a warning, unless you call .retain_grad(). Accumulation lets several micro-batches build one large-batch gradient, but forgetting to reset silently doubles every gradient on the next backward, so training loops set .grad to None or call zero_grad() each step. Updating a leaf in place raises an error, the in-place risk from the previous topic, so updates run under torch.no_grad().

Recording is wasted work whenever you will not call backward: scoring a validation set, serving a deployed model, or applying a weight update that must not itself join the graph. Inside with torch.no_grad(): results get requires_grad=False and no grad_fn, and no activations are saved. torch.inference_mode() also skips version tracking, but its tensors can never re-enter autograd: using one in a recorded computation raises Inference tensors cannot be saved for backward. detach() acts on one tensor: $\text{detach}(y)$ equals $y$ in value and shares its storage, like the views from the Tensors topic, yet contributes zero derivative to any upstream input $x$. Freezing a parameter with requires_grad_(False) is the permanent version.

Backward frees the saved tensors it consumes, so a second backward through the same graph raises Trying to backward through the graph a second time; retain_graph=True keeps them for another pass. Higher-order derivatives need more: the gradient itself must be recorded. In the formula, $y$ is an output, $x$ an input, $dy/dx$ the first derivative and $d^2y/dx^2$ its derivative, the second derivative from Mathematical Foundations. Passing create_graph=True to torch.autograd.grad or .backward() returns a gradient with its own grad_fn, ready to differentiate again. Gradient penalties, meta-learning and Hessian-vector products depend on this. Without the flag the gradient is a constant and the second call fails.

In the autograd topic you tracked every leaf tensor yourself; a real network has dozens, and losing one means it never trains. An nn.Module $M$ does the bookkeeping. $\text{own}(M)$ is the set of nn.Parameter tensors assigned directly in __init__, and $\text{children}(M)$ are the submodules assigned the same way. nn.Module.__setattr__ intercepts self.fc1 = nn.Linear(...) and files it as a child, so parameters() walks the whole tree and named_parameters() yields dotted names like fc1.weight. Call model(x), not model.forward(x), so hooks run. Failure modes: skipping super().__init__() raises an AttributeError, and a plain tensor attribute is silently ignored.

Most networks are mostly nn.Linear, so its shapes must be second nature. $\mathbf{X}$ is the $N \times d_{\text{in}}$ batch, one row per example. nn.Linear(in_features, out_features) stores weight $\mathbf{W}$ as $d_{\text{out}} \times d_{\text{in}}$, out-features first, hence the transpose, and bias $\mathbf{b}$ of length $d_{\text{out}}$, broadcast across rows as in the tensor-operations topic. Each output entry is the dot product of one input row with one row of $\mathbf{W}$, the matrix-vector product from Linear Algebra. Any leading dimensions pass through. Failure mode: if the last input dimension is not $d_{\text{in}}$, PyTorch raises mat1 and mat2 shapes cannot be multiplied.

Writing forward for a plain stack of layers is boilerplate. nn.Sequential implements the composition above: each $f_i$ is one layer, applied in order, and net[0] indexes them. Activations such as nn.ReLU hold no parameters; they are the nonlinear $f_i$ between linear ones. nn.Flatten() reshapes (B, 1, 28, 28) into (B, 784) because it starts at dimension 1 and keeps the batch, while torch.flatten starts at 0 and would merge examples. When forward needs loops or branches, store layers in nn.ModuleList, or nn.ModuleDict for named heads. Failure mode: a plain Python list of layers is invisible, so parameters() reports zero. Composition is the function chaining of Mathematical Foundations.

Some layers must behave differently while learning and while predicting, and model.train() / model.eval() flip a training flag on every submodule at once. In dropout, $p$ is the drop probability, $m_i$ a random 0-or-1 mask and $x_i$ an input; survivors are scaled by $1/(1-p)$ so the expected output still equals $x_i$. Batch normalization trains with the current batch's mean and variance while updating running_mean and running_var, which are buffers: saved state that is not learned, registered with register_buffer. In eval it uses those buffers. Forgetting eval() makes predictions random; eval() does not stop autograd, so pair it with torch.no_grad().

The scale of the random starting weights decides whether a signal survives twenty layers or fades to zero before training begins. $W_{ij}$ is one weight, $\text{fan}_{\text{in}}$ is in_features (the weight's second dimension) and $\text{fan}_{\text{out}}$ is out_features. A pre-activation sums $\text{fan}_{\text{in}}$ products, so its variance is $\text{fan}_{\text{in}} \operatorname{Var}(W)$ times the input's mean square, the variance-of-a-sum rule from Probability & Statistics. ReLU zeroes about half, hence Kaiming's factor 2; Xavier targets tanh-like activations. nn.init.kaiming_normal_ works in place. Failure modes: too small vanishes, too large explodes into NaN.

A training loop should not care whether samples live in a tensor, a folder of JPEGs or a database. A map-style dataset $D$ is that contract: $N$ is the number of samples, $x_i$ is the input of sample $i$ and $y_i$ its label. TensorDataset(X, y) implements it for tensors that share a first dimension, returning (X[i], y[i]). A custom dataset subclasses Dataset and loads lazily inside __getitem__, so memory holds one sample, not the whole corpus. The failure mode is a mismatched first dimension or an off-by-one __len__, which raises IndexError mid-epoch. It is the NumPy row indexing from NumPy & Data Wrangling, behind an interface.

Images arrive as uint8 arrays of shape $(H, W, 3)$ with values 0 to 255, but convolution layers expect float tensors of shape $(3, H, W)$ centred near zero. Here $x_{h,w,c}$ is the raw byte at row $h$, column $w$, channel $c$; $\mu_c$ and $\sigma_c$ are that channel's mean and standard deviation after scaling to $[0, 1]$; $x'$ is the result, channels first. ToTensor does the division by 255 and the HWC-to-CHW move, Normalize subtracts $\mu_c$ and divides by $\sigma_c$, and Compose chains them so __getitem__ calls one function. Failure mode: ToTensor only rescales uint8 or PIL input. This is per-channel standardisation from NumPy & Data Wrangling, applied per sample.

Gradient descent on one sample is noisy and on all $N$ samples is slow, so training uses mini-batches. $N$ is the dataset length and $B$ the batch_size; the loader yields $\lceil N/B \rceil$ batches, the last one holding the remainder $N - B\lfloor N/B \rfloor$ if it is not zero. shuffle=True installs a RandomSampler that draws a fresh permutation every epoch, while the default SequentialSampler walks $0, 1, \dots, N-1$. drop_last=True discards the short final batch. Failure modes: an unshuffled, class-sorted dataset gives single-class batches, and a final batch of size 1 crashes BatchNorm in training mode. This is the mini-batch gradient from the calculus chapter, made concrete.

After the sampler picks $B$ indices, the loader must merge a list of $B$ samples into tensors. The default collate stacks each field along a new first dimension: $B$ samples $x_1, \dots, x_B$, each of shape $(d,)$, become one tensor of shape $(B, d)$; tuples and dicts are collated field by field. Stacking needs equal shapes, so sentences of different lengths raise a RuntimeError from torch.stack. A custom collate_fn fixes this by padding every sequence to the longest in the batch with pad_sequence and returning the true lengths so the model can mask padding. It is the torch.stack from Tensor Operations, applied per batch.

With 900 negatives and 100 positives, always predicting the majority scores 90 percent, and a shuffled batch of 100 holds about 10 positives. A sampler decides which indices the loader requests. WeightedRandomSampler draws index $i$ with probability $w_i / \sum_j w_j$, where $c(i)$ is the class of sample $i$, $n_{c(i)}$ is how many samples share that class, and $w_i$ is the sample's weight. Inverse-frequency weights give each class the same total weight, so batches come out roughly balanced. It needs replacement=True to draw minority samples repeatedly, and it is mutually exclusive with shuffle=True. Repetition risks overfitting the minority. These are the categorical probabilities from Probability & Statistics.

Autograd only computes gradients: nothing changes the weights until you ask, and PyTorch never resets gradients for you. In the formula $\theta_t$ are the parameters after $t$ updates, $f_\theta(x_t)$ is the forward pass on mini-batch $x_t$ with targets $y_t$, $L$ is the scalar loss, $\nabla_\theta L$ its gradient and $\eta$ the learning rate. In code that is five lines: optimizer.zero_grad() clears .grad, pred = model(x) runs forward, loss = loss_fn(pred, y) scores it, loss.backward() fills .grad, and optimizer.step() applies the update. Because backward() adds into .grad, the accumulation rule from Autograd & the Computational Graph, skipping zero_grad() sums every past gradient and the effective step keeps growing.

A loss turns predictions into one number whose gradient points somewhere useful; the classic bug is feeding it the wrong kind of prediction. In the formula $z$ is the vector of $C$ raw scores (logits) that the last nn.Linear emits, $y$ is the integer index of the correct class and $z_y$ its score. nn.CrossEntropyLoss is exactly log_softmax plus negative log-likelihood, fused via the log-sum-exp trick so it never overflows. Pass raw logits and int64 class indices: float targets of shape (N,) raise expected scalar type Long but found Float. For binary or multi-label outputs use nn.BCEWithLogitsLoss; for regression, nn.MSELoss. This is the softmax and log-likelihood from Probability & Statistics.

Plain SGD uses one step size everywhere, crawling along shallow valleys and zig-zagging across steep ones. Momentum keeps a velocity $b \leftarrow \mu b + g$ so consistent gradients add up. Adam, shown above for parameters $\theta$ and learning rate $\eta$, keeps $\hat m$, a bias-corrected mean of gradients, and $\hat v$, a mean of their squares; $\epsilon = 10^{-8}$ guards the division, so steps are about $\eta$ at any gradient scale. The $\lambda \theta$ term is AdamW's decoupled weight decay; Adam(weight_decay=...) instead adds it to the gradient, where the $\sqrt{\hat v}$ division distorts it. Param groups give layers their own $\eta$, extending gradient descent from Mathematical Foundations.

A constant learning rate large enough for early progress keeps the weights bouncing around the minimum late in training, so a scheduler changes $\eta$ over time. In cosine annealing $\eta_t$ is the rate at step $t$, $\eta_{\max}$ the starting rate, $\eta_{\min}$ the floor (0 by default) and $T$ (T_max) the length of the decay. StepLR instead multiplies by gamma every step_size epochs. LinearLR warmup ramps up from a fraction of $\eta_{\max}$ so Adam's noisy early $\hat v$ cannot launch huge steps; SequentialLR chains warmup into cosine. Call scheduler.step() after optimizer.step(), once per unit that T_max counts; reversing the order skips the first value and PyTorch warns.

One bad batch can produce a gradient thousands of times larger than usual, flinging the weights somewhere unrecoverable and the loss to nan. Gradient clipping caps the step: $g$ is every parameter gradient concatenated into one vector, $\lVert g \rVert_2$ its Euclidean length and $c$ the max_norm threshold. torch.nn.utils.clip_grad_norm_ rescales $g$ in place between backward() and step(), keeping its direction, and returns the pre-clip norm, worth logging. Two habits finish the loop. Validate under model.eval() (dropout off, batch norm on running statistics, as in Building Models with nn.Module) and torch.no_grad(). Checkpoint optimizer and scheduler state_dict with the model's: moment buffers are training state too.

When you invent an operation, or want a different gradient than the math implies, autograd has nothing to record. A torch.autograd.Function fills that gap. Here $x$ is the input, $y = f(x)$ the output, $L$ the scalar loss, $\partial L / \partial y$ the upstream gradient (shaped like $y$), $f'(x)$ the local derivative and $\odot$ the element-wise product. forward(ctx, x) computes $y$ and calls ctx.save_for_backward(x); backward(ctx, grad_out) returns one gradient per forward input, None for non-tensors. Returning too few raises a RuntimeError; editing a saved tensor in place trips a version-counter error. It is the vector-Jacobian product from Autograd & the Computational Graph, written by hand.

Write a custom Function when the true derivative is useless or numerically broken. First case: $\hat{x} = \operatorname{round}(x)$ is a staircase whose true gradient is 0, so nothing upstream learns. The straight-through estimator (STE) keeps rounding in forward but has backward return the upstream gradient $\partial L / \partial \hat{x}$ unchanged; $:=$ marks a chosen rule, not a chain-rule result from Mathematical Foundations. The second case is stability: autograd differentiates $\log(1 + e^x)$ through exp, which overflows at $x = 100$ and returns nan, while a hand-written backward returning $\sigma(x)$ stays finite. The STE is biased, so clipped variants zero it outside $[-1, 1]$.

A wrong backward does not crash: training just runs on bad gradients and quietly underperforms. torch.autograd.gradcheck(fn, (x,)) catches this by comparing your analytic Jacobian with a numerical one. In the formula, $f_j$ is output element $j$, $x_i$ input element $i$, $e_i$ the unit vector nudging only $x_i$, and $\epsilon$ a tiny step (default $10^{-6}$). Each central difference fills one column of the numerical Jacobian; any entry differing beyond atol=1e-5 plus rtol=1e-3 raises GradcheckError. Inputs must be float64 with requires_grad=True: in float32, rounding error in the numerator swamps the difference and even a correct backward fails. This is the finite-difference derivative from Mathematical Foundations, repurposed as a test oracle.

To log what layer 12 outputs, you could edit its forward, but often you do not own that code. module.register_forward_hook(fn) attaches a callback instead. In the formula, $m$ is the module, $\text{args}$ its tuple of positional inputs and $y$ its output. A pre-hook (register_forward_pre_hook) runs before forward and may return replacement inputs; a forward hook runs after and may return a replacement output. Both fire on every call, so a module used twice fires twice. Store out.detach(): a graph-attached tensor keeps the whole forward graph alive. Each registration returns a RemovableHandle; it all relies on the submodule tree from Building Models with nn.Module.