Multilayer Networks, Data Loaders & the Training Loop

Autograd handles the calculus; this unit assembles the rest of the everyday PyTorch workflow: define a network as a module, feed it with a data loader, and train it with the canonical loop — the exact skeleton every chapter of the book reused at larger scale.

A.5 Implementing multilayer neural networks

A multilayer perceptron stacks fully-connected (Linear) layers with nonlinearities between them — here ReLU , GELU’s hard-hinged ancestor from ch4:

</> Listing A.4 — An MLP with two hidden layers

class NeuralNetwork(torch.nn.Module):
  def __init__(self, num_inputs, num_outputs):  #A
      super().__init__()
      self.layers = torch.nn.Sequential(
          # 1st hidden layer
          torch.nn.Linear(num_inputs, 30),  #B
          torch.nn.ReLU(),
          # 2nd hidden layer
          torch.nn.Linear(30, 20),
          torch.nn.ReLU(),
          # output layer
          torch.nn.Linear(20, num_outputs),
      )

  def forward(self, x):
      logits = self.layers(x)
      return logits  #C

model = NeuralNetwork(50, 3)
# NeuralNetwork(
#   (layers): Sequential(
#     (0): Linear(in_features=50, out_features=30, bias=True)
#     (1): ReLU()
#     (2): Linear(in_features=30, out_features=20, bias=True)
#     (3): ReLU()
#     (4): Linear(in_features=20, out_features=3, bias=True)
#   )
# )

#A Subclass nn.Module, build submodules in __init__, define the dataflow in forward() — the exact pattern of every class in this book, from ch3's SelfAttention_v1 to ch4's GPTModel.
#B A Linear layer computes wx + b — a weight matrix multiply plus a bias vector, a.k.a. a feedforward or fully connected layer.
#C Raw LOGITS out, no softmax — see the note below for why this is the PyTorch convention.

The 50 → 30 → 20 → 3 network: each arrow is a Linear layer (wx + b), each hidden box ends in a ReLU. The same shape-flow reading habit as ch4’s GPT — just four boxes instead of forty.

Where do its parameters live? Count them the way the book does — and sort the table to see the first layer dominate:

Layer⇅	Calculation⇅	Parameters⇅
Hidden 1 (Linear 50→30)	50·30 + 30 bias	1530	+
Hidden 2 (Linear 30→20)	30·20 + 20 bias	620	+
Output (Linear 20→3)	20·3 + 3 bias	63	+

sum(p.numel() for p in model.parameters() if p.requires_grad) → 2,213. Click rows for the arithmetic.

Logits out, softmax later — the PyTorch convention

Models return the last layer’s raw outputs (logits) without an activation because PyTorch’s common loss functions combine softmax (or sigmoid) with the negative log-likelihood in a single class — numerically more efficient and more stable. Want probabilities? Call softmax explicitly: torch.softmax(model(X), dim=1). The book’s GPT followed this convention throughout — logits everywhere, softmax only at sampling time.

Two everyday habits

torch.manual_seed(123) makes the random weight initialization reproducible (used before every model construction in the book). And for inference, wrap the call in torch.no_grad() — PyTorch skips gradient tracking, which can mean significant memory and compute savings.

A.6 Setting up efficient data loaders

The Dataset/DataLoader pair — introduced here, used in ch2, ch6 and ch7:

</> A DataLoader with the book's standard settings

train_loader = DataLoader(
  dataset=train_ds,
  batch_size=2,
  shuffle=True,
  num_workers=0,    #A
  drop_last=True    #B
)

#A num_workers parallelizes data loading/preprocessing. 0 = load in the main process — safest in Jupyter notebooks (workers can crash them via multiprocessing issues); ~4 is empirically best on many real-world setups, hardware-dependent.
#B Drop the epoch's final batch if it's undersized, keeping every training step's batch statistics comparable.

A.7 A typical training loop

The rhythm that every training chapter of this book inherits — first seen here in its smallest form, with plain SGD :

</> Listing A.9 — Neural network training in PyTorch

import torch.nn.functional as F

torch.manual_seed(123)
model = NeuralNetwork(num_inputs=2, num_outputs=2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)  #A

num_epochs = 3
for epoch in range(num_epochs):
  model.train()
  for batch_idx, (features, labels) in enumerate(train_loader):
      logits = model(features)
      loss = F.cross_entropy(logits, labels)  #B
      optimizer.zero_grad()  #C
      loss.backward()
      optimizer.step()
      ### LOGGING
      print(f"Epoch: {epoch+1:03d}/{num_epochs:03d}"
            f" | Batch {batch_idx:03d}/{len(train_loader):03d}"
            f" | Train Loss: {loss:.2f}")
  model.eval()
  # Optional model evaluation

#A Stochastic gradient descent: step directly along the gradient, scaled by lr. Chapter 5 upgraded this single line to AdamW (adaptive moments + decoupled weight decay); everything around it stayed the same.
#B cross_entropy on raw logits — the same objective from this toy 2-class net all the way to the GPT's 50,257-way next-token prediction.
#C The four-beat rhythm: zero_grad → (loss) → backward → step. Compare ch5's train_model_simple, ch6's train_classifier_simple — this loop, verbatim, with bigger models and better bookkeeping.

</> Listing A.10 — Computing prediction accuracy

def compute_accuracy(model, dataloader):
  model = model.eval()
  correct = 0.0
  total_examples = 0
  for idx, (features, labels) in enumerate(dataloader):
      with torch.no_grad():
          logits = model(features)
      predictions = torch.argmax(logits, dim=1)  #A
      compare = labels == predictions
      correct += torch.sum(compare)
      total_examples += len(compare)
  return (correct / total_examples).item()  #B

#A Argmax over logits per example, compared against the labels — ch6's calc_accuracy_loader in miniature.
#B .item() converts the final 0-dim tensor into a plain Python float.

Key idea — one workflow, every scale

Module with a forward(), loader with batches, loop with zero_grad→backward→step, accuracy by argmax: this page is the book in miniature. Chapters 2–7 changed the module (a GPT), the data (token windows, spam labels, instruction texts) and the optimizer (AdamW) — never the workflow. What remains: persistence and GPUs.

📝 Check yourself: networks, loaders & the loop

0 / 5

1.Predict: NeuralNetwork(50, 3) has layers 50→30→20→3. How many trainable parameters?
2.Why do PyTorch models conventionally return raw logits instead of softmax probabilities?
3.What does drop_last=True do in a DataLoader, and why use it in training?
4.In the training loop, what would happen if you forgot optimizer.zero_grad()?
5.This appendix trains with SGD(lr=0.5), while ch5 used AdamW. What's the relationship?