PyTorch Fundamentals — Tensors, Autograd, nn.Module

Why this session exists

The NumPy MLP exposed affine maps, activations, gradients and updates. PyTorch supplies tensor operations, reverse-mode autograd, registered modules and optimizers so you can reuse those mechanisms rather than hand-differentiate every new model. It does not validate your labels, loss, split or intended computation for you.

You will build a CPU model with nn.Module, inspect tensor layout/dtype/device, write a DataLoader training loop, distinguish evaluation mode from gradient mode, save/load a state dictionary and match a NumPy gradient and SGD update. GPU placement is explained as an API contract; all execution in this lesson is explicitly CPU-only.

Prerequisites

Use S097 forward matrices, S098 derivatives and S099 optimizers. Know row-batch shapes and that nn.Linear stores the transpose of our NumPy weight convention.

(a) Intuition · 5 min

After building a bicycle, using a manufactured one is convenient without making its gears mysterious. Likewise, nn.Linear, .backward() and optimizer.step() replace different pieces of the scratch implementation. The analogy is about understanding and reuse, not a claim that PyTorch is always faster or every production project must use it.

The three abstractions to internalise

• Tensor: data interpreted by shape, strides, dtype, storage offset and device. It may also participate in an autograd graph.
• Autograd: records differentiable operations in grad mode, composes local derivatives and accumulates requested leaf gradients. It does not change parameters by itself.
• Module: organizes registered parameters, child modules and buffers and defines forward. Optimizers consume selected parameters; DataLoader handles sample iteration and batching separately.

The useful model is buffer + metadata + optional differentiation graph. A device mismatch concerns storage placement. A transpose/view failure concerns metadata/layout. Missing gradients or retained computation concerns the graph. Diagnose the layer of the abstraction that actually failed.

PyTorch's history in 6 events

The relevant design progression is from tensor operations, through dynamic reverse-mode differentiation and reusable modules, to optional compilation/export mechanisms. In eager PyTorch, the graph is reconstructed from the operations actually executed each iteration; Python control flow can therefore select different paths. The 2019 PyTorch design paper explicitly describes an imperative Pythonic style that treats code as the model, supporting ordinary debugging and user control. The PyTorch 2.0 release announcement then presents torch.compile as an optional wrapper while preserving eager-mode development. Its documented components separate graph capture (TorchDynamo), ahead-of-time backward tracing (AOTAutograd), primitive-operator canonicalization (PrimTorch) and code generation (TorchInductor). Compilation is therefore an additional execution path with its own constraints and overhead, not a replacement for understanding tensors, modules or gradients. This lesson makes no adoption-percentage, historical speedup or compilation benchmark claim.

The PyTorch training loop, always the same shape

The central five lines are zero_grad, model call, loss, backward, step. Clearing after a completed step is also valid if done before the next independent accumulation. Deliberate accumulation clears once per intended update, not before every microbatch. Accidentally accumulating gradients does not literally multiply the learning rate by the batch count: gradient directions and parameter values change along the way.

The numpy → PyTorch mapping

PyTorch 2.9 nn.Linear initializes weight and bias uniformly in $[-1/\sqrt{fan_{in}},1/\sqrt{fan_{in}}]$. Its implementation calls kaiming_uniform_ with a=sqrt(5) to achieve this bound; that does not mean it automatically uses ReLU He-normal variance $2/fan_{in}$. If you want that recipe, explicitly call nn.init.kaiming_normal_(weight, nonlinearity='relu') on the intended hidden layers and choose bias/output initialization deliberately.

Device management — make placement explicit

Choose torch.device('cpu') for this lab. A CUDA-capable installation can test torch.cuda.is_available() before selecting CUDA; MPS has its own torch.backends.mps.is_available() check. Merely running on a Mac does not prove MPS availability. Move the model before constructing the optimizer and move each input/target batch consistently. Prefer new tensors constructed with the intended device and dtype, or x.new_zeros(...) when they should match x.

Module.to(...) moves registered parameters and buffers recursively, not arbitrary tensors hidden in ordinary attributes or lists. Most tensor operations require compatible devices/dtypes, though specific scalar operations have exceptions; inspect the failing operation rather than assuming every possible mixed-device expression is illegal. The DataLoader returns whatever the dataset/collation produces; CPU tensors are the usual pattern, not an unconditional framework law.

Float32 is a practical training default; float64 is useful for numerical checks. Float16 has a much smaller exponent range than float32. Bfloat16 retains float32's exponent range but fewer significand bits. Supported kernels and mixed-precision policies depend on device and operation. Integer class labels cannot carry ordinary floating autograd gradients; cross entropy class indices here use torch.long and lie between zero and $K-1$.

nn.Module — the pattern for every model

Call super().__init__() before assigning child modules or parameters; assigning them before parent initialization can raise an error. self.layer=nn.Linear(...) registers a child module; nn.Parameter registers a parameter when assigned; register_buffer records state that should move and often persist without being an optimizer parameter. Use ModuleList for a list of submodules rather than an ordinary list if registration is needed.

model.parameters() yields registered parameters, including those with requires_grad=False. Freezing parameters is distinct from changing train/eval mode. A plain tensor attribute is not registered merely because it has requires_grad=True.

Call model(x) rather than directly calling model.forward(x) so module hooks and call machinery run. Direct forward calls do not inherently disable autograd; the differentiable operations inside still record in grad mode. model.train() and model.eval() recursively set the training flag for mode-sensitive layers such as dropout and BatchNorm. A model made only of Linear/ReLU has the same mathematical forward function in both modes.

Autograd, leaves and logging

A differentiable operation is recorded in default grad mode when at least one relevant input requires gradients. requires_grad=True does not mean every subsequent operation is differentiable; integer conversion, detach, Python scalar extraction and explicit no-grad regions can break the path. Leaf tensors normally accumulate into .grad; nonleaf tensors need retain_grad() if you want their .grad stored for inspection. A non-scalar output requires a supplied cotangent to .backward(...) or a scalar reduction.

model.eval() does not turn autograd off. Use torch.no_grad() for evaluation without a gradient graph, or inference_mode() when its stronger restrictions are suitable. Neither is a blanket guarantee of a particular speedup. Some analyses intentionally need input gradients while the model remains in evaluation mode.

Log scalar loss with .item() or keep detached tensor summaries. Appending live losses can retain graph objects and, depending on backward/retention state, saved tensors. Saved intermediates are normally freed by backward, so “every appended loss retains the full activation graph forever” is too strong. Still, retaining unnecessary live graph references is wasteful and may keep larger structures alive. .detach() shares storage; use .detach().clone() if an independent snapshot is required.

Views, strides and copies

For an ordinary strided tensor, an element's address is determined by storage offset plus the dot product of indices and strides. Transposing usually changes metadata rather than moving all elements. Basic slicing similarly returns views; advanced indexing can copy. Views save data-buffer copies but still create metadata objects and have nonzero overhead.

A contiguous 2×3 tensor has strides (3,1); its transpose has shape 3×2 and strides (1,3). Flattening the transpose with .view(-1) cannot express its requested logical order as that simple compatible view and raises. .reshape(-1) can copy to satisfy the request. .contiguous() returns the original tensor if already in the desired contiguous format and otherwise copies. Do not depend on whether reshape happens to alias; if independence matters, clone explicitly. Noncontiguous operations are not universally slower—kernels may support their layouts efficiently.

Where should preprocessing live?

Dataset/CPU transforms are flexible and may use workers, but workers add process, memory and transfer overhead; num_workers=0 is a useful small-CPU baseline. Tensor-expressible transforms can run with the model where that is supported, but consume its compute/memory. Offline deterministic preprocessing avoids repetition but requires versioned artifacts and regeneration when the recipe changes.

Fit statistical preprocessing on training rows only, even when precomputing offline. Stochastic augmentation should respect label invariances and can be sampled online; a finite cached set of random variants is also possible, with less diversity, not categorically forbidden. Use profiling to identify an input bottleneck before changing workers or device placement. These trade-offs are conceptual here; no accelerator utilization or throughput measurement was made.

(c) Hands-on · 25 min

The default uses bundled digits, CPU and zero worker subprocesses. It learns normalization from training only, stores it as module buffers, monitors validation, restores a cloned best state and evaluates the test once. Save/load is demonstrated with an in-memory buffer, so running the snippet creates no persistent checkpoint file.

import copy
import io
import numpy as np
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
 
torch.set_num_threads(1); torch.manual_seed(0)
DEVICE=torch.device('cpu')
class MLP(nn.Module):
    def __init__(self,in_dim=64,h1=32,h2=16,out_dim=10):
        super().__init__()
        self.register_buffer('mean',torch.zeros(in_dim))
        self.register_buffer('scale',torch.ones(in_dim))
        self.fc1=nn.Linear(in_dim,h1); self.fc2=nn.Linear(h1,h2); self.fc3=nn.Linear(h2,out_dim)
    def forward(self,x):
        x=x.flatten(start_dim=1)
        x=(x-self.mean)/self.scale
        return self.fc3(torch.relu(self.fc2(torch.relu(self.fc1(x)))))
 
X,y=load_digits(return_X_y=True); X=torch.tensor(X,dtype=torch.float32); y=torch.tensor(y,dtype=torch.long)
dev,test=train_test_split(np.arange(len(y)),test_size=.2,stratify=y.numpy(),random_state=1)
train,val=train_test_split(dev,test_size=.25,stratify=y[dev].numpy(),random_state=2)
assert not (set(train)&set(val) or set(dev)&set(test))
def loader(ids,shuffle=False):
    return DataLoader(TensorDataset(X[ids],y[ids]),batch_size=64,shuffle=shuffle,
        num_workers=0,generator=torch.Generator().manual_seed(3))
m=MLP().to(DEVICE)
with torch.no_grad():
    m.mean.copy_(X[train].mean(0)); m.scale.copy_(X[train].std(0,unbiased=False).clamp_min(1.))
opt=torch.optim.Adam(m.parameters(),lr=.01); criterion=nn.CrossEntropyLoss()
tr_loader=loader(train,True)
def run(model,ld,optimizer=None):
    training=optimizer is not None; model.train(training)
    total=0.; correct=0; count=0
    with torch.set_grad_enabled(training):
        for x,yb in ld:
            x,yb=x.to(DEVICE),yb.to(DEVICE)
            assert x.dtype==torch.float32 and yb.dtype==torch.long
            if training: optimizer.zero_grad(set_to_none=True)
            z=model(x); loss=criterion(z,yb)
            assert torch.isfinite(loss)
            if training:
                loss.backward()
                assert all(p.grad is None or torch.isfinite(p.grad).all() for p in model.parameters())
                optimizer.step()
            total+=loss.item()*len(yb); correct+=(z.argmax(1)==yb).sum().item(); count+=len(yb)
    return total/count,correct/count
best=float('inf')
for epoch in range(1,16):
    tr=run(m,tr_loader,opt); va=run(m,loader(val))
    if va[0]<best: best=va[0]; state=copy.deepcopy(m.state_dict()); chosen=epoch
    if epoch in [1,5,15]:print(epoch,'online train',tr,'validation',va)
m.load_state_dict(state); m.eval()
print('chosen',chosen,'eval-mode training',run(m,loader(train)),'test',run(m,loader(test)))
# Includes fitted buffers, not only trainable weights.
buf=io.BytesIO(); torch.save(m.state_dict(),buf); buf.seek(0)
restored=MLP(); restored.load_state_dict(torch.load(buf,map_location='cpu',weights_only=True)); restored.eval()
with torch.no_grad(): assert torch.equal(m(X[:8]),restored(X[:8]))
assert 'mean' in restored.state_dict() and 'mean' not in dict(restored.named_parameters())
large=MLP(784,128,64,10)
assert sum(p.numel() for p in large.parameters())==109386
assert large(torch.zeros(2,1,28,28)).shape==(2,10)
print('state/buffer roundtrip and 784-input contracts passed')
# Storage/layout tests.
a=torch.arange(6.).reshape(2,3); t=a.T
assert t.stride()==(1,3)
t[0,1]=99; assert a[1,0]==99
try: t.view(-1)
except RuntimeError: pass
else: raise AssertionError('expected incompatible view')
r=t.reshape(-1); assert r.shape==(6,) and torch.equal(r,torch.tensor([0.,99.,1.,4.,2.,5.]))
x=torch.tensor(2.,requires_grad=True); detached=x.detach()
assert detached.data_ptr()==x.data_ptr() and not detached.requires_grad
(3*x).backward(); (3*x).backward(); assert x.grad.item()==6 and x.item()==2
m.eval(); assert m(X[:2]).requires_grad
with torch.no_grad(): assert not m(X[:2]).requires_grad
print('views, detach, accumulation and mode separation passed')

Anatomy of the script

The training and validation loaders never touch test indices. The scaler buffers move with the model and survive the state roundtrip. Batch-mean losses are multiplied by actual batch size before summing so a final short batch is weighted correctly. A copied state_dict is required: merely assigning best=model.state_dict() can retain references to subsequently changing tensors.

To adapt the same model to trusted official MNIST arrays, keep its 60,000-image training and 10,000-image test roles intact, carve validation only from training, supply 28×28 images and instantiate MLP(784,128,64,10). Fit normalization only on the chosen training subset. The displayed 784-input fixture verifies flattening and parameter count but is not a MNIST training run. No mixed-precision support is claimed for a loop that does not contain it.

Practice: match the NumPy backward and update

Copy using with torch.no_grad(): layer.weight.copy_(torch.from_numpy(W.T)), not .data reassignment. Match dtype, activation, loss reduction and optimizer before comparing. The following standalone float64 fixture checks both weights and bias gradients and one SGD update; S098 checks the whole multilayer chain.

import numpy as np
import torch
from torch import nn
from scipy.special import logsumexp
torch.set_num_threads(1)
rng=np.random.default_rng(12)
X=rng.normal(size=(5,3)); W=rng.normal(size=(3,2)); b=rng.normal(size=2); y=np.array([0,1,1,0,1])
z=X@W+b; lp=z-logsumexp(z,axis=1,keepdims=True)
dz=np.exp(lp); dz[np.arange(5),y]-=1; dz/=5
dW=X.T@dz; db=dz.sum(0)
layer=nn.Linear(3,2).double()
with torch.no_grad(): layer.weight.copy_(torch.from_numpy(W.T)); layer.bias.copy_(torch.from_numpy(b))
optimizer=torch.optim.SGD(layer.parameters(),lr=.1)
optimizer.zero_grad(set_to_none=True)
loss=nn.functional.cross_entropy(layer(torch.from_numpy(X)),torch.from_numpy(y)); loss.backward()
assert np.allclose(layer.weight.grad.numpy().T,dW,atol=1e-12)
assert np.allclose(layer.bias.grad.numpy(),db,atol=1e-12)
optimizer.step()
assert np.allclose(layer.weight.detach().numpy().T,W-.1*dW,atol=1e-12)
assert np.allclose(layer.bias.detach().numpy(),b-.1*db,atol=1e-12)
print('NumPy/nn.Linear gradient and SGD step parity passed')

(d) Production reality · 15 min

Gradient lifecycle: Test that backward leaves parameters unchanged and doubles a repeated fresh graph's identical gradient. Then test that a step changes the expected registered parameters. A None gradient can mean no graph path or a cleared buffer, not necessarily a numerical zero. set_to_none=True can also change optimizer behavior compared with a zero tensor for unused parameters.

Inference lifecycle: Load trusted compatible state, move model, set mode deliberately and choose grad/no-grad/inference mode for the task. Missing eval changes dropout and default BatchNorm behavior, but a plain Linear/ReLU network need not change accuracy. Avoid invented 10× latency anecdotes; measure the real serving path when authorized.

Persistence: State dictionaries include parameters and persistent buffers, not architecture code, feature meanings or all training state. Save the class/config, preprocessing schema, label mapping and dependency versions alongside them for a real artifact. For resumption also save optimizer/scheduler state and RNG/counters as required. weights_only=True restricts deserialization; do not treat arbitrary untrusted checkpoints as safe merely because they are called weights. Whole-object saving uses pickle and is coupled to Python classes; it does not simply embed a universally portable class definition.

Reproducibility: Seed each RNG used, specify data order/workers and record versions/device. torch.manual_seed alone does not seed a separate NumPy Generator. Deterministic-algorithm settings can raise for unsupported operations or change performance; even careful settings do not promise bitwise reproducibility across releases/platforms. The lab's seeds make its local tests inspectable, not universal hardware parity.

Where this shows up in the rest of the plan

S101 relies on train/eval state and saved buffers. CNNs, recurrent modules and attention use the same registration and gradient lifecycle but different forward operations. Export and compilation are separate compatibility projects, not automatic consequences of a successful state-dictionary roundtrip.

(e) Recall + stretch · 10 min

1. What does requires_grad do? Requests differentiation through supported operations in grad mode; leaf gradients accumulate when connected to the backward output.
2. What does backward do? Propagates the supplied or scalar-unit cotangent and accumulates leaf gradients, without an optimizer step.
3. Why clear gradients? To begin a new intended accumulation group, not necessarily every microbatch.
4. Five lines? Clear, model call, loss, backward, step. Logging uses detached values.
5. Train versus eval? Changes module mode; no-grad changes graph recording. They are independent.
6. Why logit CE? Stable log-sum-exp avoids taking logs of underflowed probabilities. Targets and reduction still need checking.
7. Device error? Inspect registered state and all operands; move them consistently and verify availability before choosing an accelerator.
8. When no-grad? When those derivatives are not needed. Eval-mode saliency or input optimization may intentionally keep gradients enabled.
9. Why state_dict? Explicit parameters/buffers separated from architecture code, with compatible configuration still required.
10. Stretch: Transpose copied NumPy weights, match float64 mean CE and SGD, and compare every gradient/update. An autograd comparison is a strong independent check, not proof that the intended objective or data is correct.

Explain-out-loud test

Name the five training actions, explain buffer versus parameter versus ordinary attribute, and diagnose one storage, one layout and one graph error without conflating their fixes.

What comes next

Dropout, normalization and decay introduce stateful/stochastic behavior that makes these distinctions operationally important.

Sources and execution boundary

PyTorch 2.9 autograd, modules, tagged Linear implementation, tensor views, reproducibility, and cross entropy. Both displayed programs are CPU-executed. No GPU/MPS, compilation/export, MNIST download or serving benchmark is asserted.