Building a Network in PyTorch

PyTorch’s training loop consumes batches, and DataLoader is the component that turns raw arrays into batches: it slices, shuffles, casts, and optionally prefetches in worker processes. The contract any dataset must satisfy is two methods — len() and getitem(i) — and for in-memory arrays you never have to write them yourself:

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

X = torch.tensor(features, dtype=torch.float32)              # (N, 20)
y = torch.tensor(targets, dtype=torch.float32).unsqueeze(1)  # (N, 1)

ds = TensorDataset(X, y)
train_ds, val_ds = random_split(ds, [0.8, 0.2],
    generator=torch.Generator().manual_seed(42))

train_dl = DataLoader(train_ds, batch_size=64, shuffle=True)
val_dl = DataLoader(val_ds, batch_size=128)        # no shuffle for evaluation

for xb, yb in train_dl:
    print(xb.shape, yb.shape)                       # torch.Size([64, 20]) torch.Size([64, 1])
    break

Data that does not fit in one tensor — images on disk, variable-length records — means writing the class yourself:

class MyDataset(torch.utils.data.Dataset):
    def __init__(self, records, transform=None):
        self.records = records
        self.transform = transform

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

    def __getitem__(self, i):
        rec = self.records[i]
        x = self.transform(rec['signal']) if self.transform else rec['signal']
        return torch.tensor(x, dtype=torch.float32), rec['label']

The default collate function stacks the per-item tensors into a batch dimension — that is where (B, ...) appears. For ragged sequences you replace collate_fn, usually by padding to the longest item in the batch.

Check the first batch’s shapes before training anything. A (64, 1) target fed to CrossEntropyLoss, or a float where the loss wants Long, fails on step one — far cheaper than 30 epochs later. Then get the loop right:

model = MLP(n_in=20, n_hidden=64, n_out=1)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = torch.nn.MSELoss()

for epoch in range(30):
    model.train()
    total, count = 0.0, 0
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_fn(pred, yb)
        opt.zero_grad()
        loss.backward()
        opt.step()
        total += loss.item() * xb.shape[0]     # weight by batch size
        count += xb.shape[0]
    print(f'epoch {epoch} train_mse={total / count:.4f}')

Two settings carry most of the decision weight. shuffle=True belongs on the training loader only: SGD assumes each mini-batch gradient is an independent estimate of the population gradient, and a fixed ordering feeds correlated batches that let the model settle into oscillations tracking the data order. Evaluation loaders stay unshuffled — order does not matter there, and determinism makes runs comparable. And batch size: 32 to 256 is the usual tabular range, larger batches generally want a scaled learning rate, and drop_last=True avoids a final one-sample batch that wrecks BatchNorm statistics — a preview of the mode-sensitive layers in the next chapter.

Two operational knobs matter once the loop runs. num_workers spawns subprocesses that load and collate the next batch while the GPU trains on the current one; on Windows, where spawning is the process start method, keep DataLoader construction and the training loop inside an if name == 'main': guard or the workers fail before they start. persistent_workers=True keeps those workers alive across epochs instead of respawning them per pass — a real speedup when data loading is the bottleneck. Neither knob changes what the model sees; they change how long the GPU waits for it. And note what random_split did back at the top: it slices one dataset into index views rather than copying tensors, so memory stays flat no matter how you divide it.

One habit from the very first run: decide the device once and apply it everywhere, which is the third chapter’s subject. The middle three lines of the loop — zero_grad, backward, step — are dissected in Backpropagation and the Training Loop; everything from here on assumes that loop is muscle memory.

An nn.Module is a container PyTorch can inspect. Subclass it and assign the layers in init:

class MLP(torch.nn.Module):
    def __init__(self, n_in, n_hidden, n_out):
        super().__init__()
        self.fc1 = torch.nn.Linear(n_in, n_hidden)
        self.fc2 = torch.nn.Linear(n_hidden, n_out)

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

model = MLP(20, 64, 3)
print(sum(p.numel() for p in model.parameters()))   # 1539 = (20·64+64) + (64·3+3)

Assignment is registration. Putting an nn.Linear on self inserts it into the module’s internal _modules dict, so model.parameters() recurses into it, the optimizer sees it, and .to(device), .train(), and state_dict() all reach it. A raw torch.nn.Parameter stored on self is registered too. But a plain tensor — self.scale = torch.ones(3) — is invisible machinery: it never appears in parameters(), receives no gradients, never moves with .to(device), and trains nothing. The model builds, the forward pass runs, the loss decreases, and your scale stays exactly 1. The reveal is one line: state_dict() lists registered tensors and omits the impostor. When a tensor is meant to be learned, use nn.Parameter; when it is meant to be saved but not trained — a counter, a mask — register_buffer('name', tensor) puts it in state_dict without making it trainable.

The arithmetic from Perceptron and Multilayer Networks still audits the module: (20·64 + 64) + (64·3 + 3) = 1,344 + 195 = 1,539 trainable parameters, and named_parameters() prints each name exactly as the checkpoint will store it — a ten-second pre-training check that every intended layer registered and no bare tensor is along for the ride.

Calling the model is another place where the wrapper matters. model(x) invokes call, which runs pre-hooks and post-hooks around your forward() — the mechanism mixed-precision wrappers, profilers, gradient checkpointing, and quantization hooks all rely on. Calling model.forward(x) directly bypasses that machinery and silently disables anything that instrumented the module. Framework code, and other people’s models, will assume you called the module: always call model(x).

Then the two flags interviewers know to ask about. model.train() does not mean start training, and model.eval() does not mean stop — neither touches gradients or the optimizer. They set a mode flag that specific submodules read:

Mis-set the flags and the symptoms are specific. Forget model.eval() before validation and your metrics fluctuate between runs on identical weights and data — dropout is randomly deleting units while you measure, and BatchNorm mixes each batch’s statistics into the score. Forget to switch back to model.train() after evaluating and dropout stays active during training: convergence slows in a way that looks exactly like a learning-rate problem but is not. Neither mode flag touches autograd — that is torch.no_grad(), a separate mechanism:

model.eval()
with torch.no_grad():
    for xb, yb in val_dl:
        pred = model(xb)          # no graph retained: less memory, faster

Train and eval decide what the layers compute; no_grad decides whether the graph is recorded at all. Use all three, in that combination, every time you evaluate. For fine-tuning, the related lever is requires_grad_(False) on the backbone: frozen weights receive no gradients, appear as non-trainable in the optimizer’s parameter groups, and skip their share of the backward cost — which is how a fine-tune of a large model fits on hardware that could never retrain it from scratch. The next step is making the whole workflow survive a crash — checkpoints and devices — which is also where the pieces assemble into the pipeline you will reuse for Convolution and Pooling.

model.state_dict() returns an OrderedDict mapping names to tensors: fc1.weight as (64, 20), fc1.bias as (64,), fc2.weight, fc2.bias — and any registered buffers, such as BatchNorm’s running_mean, running_var, and num_batches_tracked. Buffers matter more than their quiet name suggests: they are state the model uses at inference but never trains, so a checkpoint that loses them hands your validation pass wrong statistics with no error at all. Save the dict, and load it into an identical architecture:

torch.save(model.state_dict(), 'best.pt')

clone = MLP(20, 64, 3)                    # architecture must exist before loading
clone.load_state_dict(torch.load('best.pt', weights_only=True))

Save state_dict, not the module. torch.save(model) pickles the entire object including the class’s import path and code version — it breaks the moment the class moves, gets renamed, or the PyTorch version changes, and unpickling files from the internet executes arbitrary code, which is exactly why weights_only=True exists and should be your default. A flat name-to-tensor map has none of those problems.

Resuming interrupted training, though, needs more than weights. Adam and AdamW keep per-parameter moment estimates in the optimizer — momentum accumulated over thousands of steps — and a fresh optimizer starts those at zero, so a resumed run loses its accumulated direction and the loss bumps before recovering. Persist all of it together:

torch.save({
    'epoch': 37,
    'model': model.state_dict(),
    'opt': opt.state_dict(),
    'best_val': 0.812,
}, 'ckpt.pt')

ckpt = torch.load('ckpt.pt', weights_only=False, map_location=device)
model.load_state_dict(ckpt['model'])
opt.load_state_dict(ckpt['opt'])
start_epoch = ckpt['epoch'] + 1

Note the weights_only exception: optimizer state_dicts contain more than plain tensors, so resuming needs weights_only=False — accept that only for checkpoints you produced yourself.

load_state_dict defaults to strict=True, which errors when the checkpoint’s keys and the live architecture disagree — a missing head or a renamed layer fails loudly instead of silently leaving those parameters at random initialization. And after any load, run one validation batch through the restored model and compare its loss against the recorded best_val: that single check confirms weights, mode flags, and preprocessing all line up before you trust a resumed run’s numbers.

Which checkpoint should exist at the end is a separate decision, and the default is wrong. The last epoch is rarely the best one, so track the validation metric and save on improvement:

best = float('inf')
for epoch in range(epochs):
    train_one_epoch(model, train_dl, opt, loss_fn)
    val = evaluate(model, val_dl)              # eval() + no_grad, from chapter two
    if val < best:
        best = val
        torch.save(model.state_dict(), 'best.pt')

Report results from the best checkpoint, not the final weights — otherwise your quoted number is a coin flip on where training happened to end.

Last, device discipline, the most common runtime error in PyTorch code. The model moves once, every batch moves every loop:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

for xb, yb in train_dl:
    xb, yb = xb.to(device), yb.to(device)     # forgetting this line raises:
    # Expected all tensors to be on the same device, but found at least two devices

The error is honest: the model’s weights live on the GPU and the batch arrived on the CPU, or vice versa. And when loading a checkpoint written on a GPU machine onto a CPU-only one, torch.load(..., map_location='cpu') remaps the stored device tags instead of failing to allocate CUDA memory that does not exist. With datasets feeding batches, modules owning parameters, the paired loss from Output Layers and Loss Pairings getting backward()’d by the loop in Backpropagation and the Training Loop, and checkpoints preserving all of it, the workflow is complete — the same skeleton runs the convolutional models in Convolution and Pooling.