Take Lesson 3's hand-rolled loop and swap each piece for the library version: nn.Linear holds the parameters, nn.MSELoss is the loss, torch.optim takes the step. Same five moves, same result — now written exactly the way real code is.
Real PyTorch never hand-manages w, b, and w -= lr*w.grad. It wraps the parameters in a nn.Module (the model), measures error with a ready-made loss function, and lets a torch.optim optimizer do the update for all parameters at once. The five moves are unchanged — they just become five library calls. Learn this mapping and you can read almost any training script.
You already know the loop. This lesson is pure refactor: each thing you wrote by hand has a standard replacement, and that replacement is what scales from a line to a 100-layer network without changing the loop.
| Lesson 3 (from scratch) | → Lesson 4 (idiomatic) | what it is |
|---|---|---|
w, b, X @ w + b | nn.Linear(1, 1) | a module: holds params + a forward pass |
((pred-y)**2).mean() | nn.MSELoss() | a loss function object |
w -= lr*w.grad | optimizer.step() | an optimizer: updates all params |
w.grad.zero_() | optimizer.zero_grad() | clears all grads at once |
nn.Linear is just w·x + bA layer is an nn.Module that owns some parameters and knows how to transform an input. nn.Linear(1, 1) is precisely our line: one input feature → one output, with a learnable weight and bias (initialized to small random values for you).
import torch.nn as nn model = nn.Linear(in_features=1, out_features=1) print(model) # Linear(in_features=1, out_features=1, bias=True) for name, p in model.named_parameters(): print(name, p.data.flatten().tolist()) # weight [-0.2489] ← random init, requires_grad=True # bias [ 0.0451]
.forward()
You run the model with model(X), not model.forward(X). Calling the instance triggers __call__, which runs forward plus PyTorch's bookkeeping (hooks, train/eval state). Calling .forward() directly skips that — a subtle bug you'll see people make.Source: torch.nn.Linear · Build the model tutorial.
The loss is now an object you call. The optimizer is handed the model's parameters once, up front — after that, it knows what to update.
loss_fn = nn.MSELoss() # mean squared error, as before optimizer = torch.optim.SGD(model.parameters(), lr=0.1) # SGD = plain gradient descent
model.parameters() hands the optimizer every learnable tensor in the model (here: weight and bias). That's the wiring that lets one optimizer.step() update all of them — whether there are 2 parameters or 2 billion.
step() updates them all.Here is the loop you'll see everywhere. Compare it line-for-line with Lesson 3 — same five moves, now five calls (convention puts zero_grad first):
for epoch in range(201): optimizer.zero_grad() # clear old gradients pred = model(X) # 1. forward loss = loss_fn(pred, y) # 2. loss loss.backward() # 3. backward optimizer.step() # 4. step — updates ALL params at once
epoch 0 | loss 10.98034 epoch 50 | loss 0.01200 epoch 100 | loss 0.01050 epoch 200 | loss 0.01049 learned w = 2.000 (true 2.0), b = -2.996 (true -3.0)
.grad or wrote an update rule. Add 50 layers and this loop is byte-for-byte the same; only model changes. That is why everyone writes it this way.~/projects/learn/public/courses/pytorch/practice/.venv/bin/python ~/projects/learn/public/courses/pytorch/practice/lesson04_nn_module.pynn.Module — the pattern you'll read everywhereReal models aren't a bare nn.Linear; they're a subclass of nn.Module that wires several layers together. The pattern is always the same two methods:
class LinearModel(nn.Module): def __init__(self): super().__init__() # always call this first self.fc = nn.Linear(1, 1) # layers as attributes → auto-registered as params def forward(self, x): # define the computation; PyTorch calls it via model(x) return self.fc(x) m = LinearModel() print(m.state_dict().keys()) # odict_keys(['fc.weight', 'fc.bias'])
self.<name> in __init__ — PyTorch auto-discovers their parameters (that's how model.parameters() and the optimizer find them). (2) Put the data flow in forward. The state_dict is just a dict of all those named tensors — which is exactly what you save and load.Source: torch.nn.Module.
The optimizer is a one-line swap. Plain SGD works, but most modern code reaches for Adam, which adapts the step size per parameter and usually trains faster with less tuning:
optimizer = torch.optim.SGD(model.parameters(), lr=0.1) # simple, needs tuned lr optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) # adaptive; great default
The loop doesn't change at all — only the object you construct. When you read a repo, the optimizer choice and its lr are two of the first hyperparameters to spot.
model.parameters()
An optimizer constructed without the model's parameters (or with the wrong model's) will happily run and update nothing. If training does nothing, check that line.Source: torch.optim.
optimizer.step() do?nn.Module, why assign layers as self.fc = nn.Linear(...)?state_dict?nn.Linear(in, out) do?x·Wᵀ + b — a learnable weight matrix and bias. nn.Linear(1,1) is exactly the line w·x + b, with params auto-initialized and tracked.model(X) or model.forward(X)?model(X). It invokes __call__, which runs forward plus hooks and train/eval handling. Calling .forward() directly skips that.model.parameters() and a learning rate. step() updates all params from their grads; zero_grad() clears all grads.nn.Module?__init__ (call super().__init__(), then assign layers as attributes) and forward(self, x) (the computation). PyTorch runs forward when you call the instance.SGD is plain gradient descent (needs a tuned lr); Adam adapts per-parameter step sizes and is the usual default (often lr=1e-3).zero_grad() (reset) → model(X) (forward) → loss_fn(pred,y) (loss) → loss.backward() (backward) → optimizer.step() (step).Good asks: "add a second layer and a nonlinearity to the model," "what exactly is in model.parameters() for a 2-layer net?", or "show SGD vs Adam converging side by side." Next: real data doesn't arrive as one big tensor — we batch it. Say "go" for Lesson 5.