nn.Module and optim - the Honest Refactor

Context

nn.Linear is Lesson 4.2.1's neuron as an object that owns its weight and bias, nn.MSELoss is its loss line as an object, and torch.optim.SGD is its step and its zeroing held in one object with the learning rate. The three replace the lines you wrote by hand, and they compute the same numbers, which this lesson proves from the same seeded start rather than asserts.

Every training loop you will read from here on is written in them. The network that classifies which listings sell within a month in Section 4.3, the image model of Section 4.5 and any PyTorch training script you open in someone else's repository are the same five lines around these three objects: a forward pass, a loss, backward(), optimizer.step() and optimizer.zero_grad(). A model saved in Lesson 4.4.3 is the dictionary of parameters a module keeps, and a loss swapped in Lesson 4.3.2 is one object replaced.

What breaks is a refactor taken on faith. optimizer.step() read as a black box hides the step Lesson 4.1.4 took inside no_grad, a target of the wrong shape under nn.MSELoss warns once and then fits a different loss, and the zeroing left out of the refactored loop bounces exactly as it did in Lesson 4.2.1, with nothing raised. That lesson's loop settled at 144,880 and 380,326 with a loss of 1.7226e+10. This one rebuilds the loop in the three classes, piece by piece, and shows the hand loop and the module loop print the same digits at every epoch from one seeded start.

Goal

• Build nn.Linear(1, 1) under a seed, read its two nn.Parameter tensors by name and shape, run it on a (n, 1) input, and name the two refusals a one-dimensional or a float64 input meets.
• Compute the loss with nn.MSELoss on a (n, 1) target, and read the warning a (n,) target raises and the number it computes on a fitted model.
• Replace Lesson 4.2.1's step and zeroing with torch.optim.SGD, and show through torch.equal and torch.allclose what step() and zero_grad() do to each parameter.
• Run the hand loop and the module loop from the same seeded start and prove they print the same loss at every epoch and reach the same parameters.
• Write an nn.Module subclass with super().__init__() and forward, read its named parameters and its training flag, and diagnose the two refusals a subclass can raise.

Prerequisites

None of this is required, the lesson stands on its own. Both are here if you want PyTorch's own picture of a module before you build one.

• YouTube - "Building Models with PyTorch" by PyTorch. Fourteen minutes from the PyTorch channel: the Module class as the parent of every model, the Parameter class as the tensor a module learns, and a tour of the common layer types, with nn.Linear among them.
• Official docs - PyTorch Tutorials, "Build the Neural Network". Its sections "Define the Class", "Model Layers" and "Model Parameters": a subclass of nn.Module with __init__ and forward, nn.Linear applied to a batch of inputs, and named_parameters() read back, on an image model built from layers later sections introduce.

Table of Contents

• nn.Linear Is the Neuron
• nn.MSELoss Is the Loss, and the Shape It Wants
• optim.SGD Is the Step and the Zeroing
• The Same Numbers, Proved
• Your Own nn.Module

Theory

Every Theory block runs in ~/dev/bootcamp on the 2,875 listings with an area, Lesson 4.2.1's rows, as float32 tensors, with the area standardized by its own mean and std(correction=0), 99.02 and 44.0589, exactly as 4.2.1 standardized it. A Python block that opens with its imports is a script of its own, and a block that does not continues the script above it.

nn.Linear Is the Neuron

What. A module (an nn.Module, PyTorch's container for the parameters of a model and the forward that uses them) is what every layer and every model is. nn.Linear(in_features, out_features) is Lesson 4.2.1's neuron as a module: it holds a weight of shape (out_features, in_features) and a bias of shape (out_features,), each a parameter (an nn.Parameter, a tensor a module owns, created with requires_grad=True and listed by parameters()), and model(X) computes X @ weight.T + bias for an input of shape (n, in_features). Why. Lesson 4.2.1's w and b were two names in a script. A network has layers of them, and a layer that owns its parameters hands all of them to an optimizer in one call, prints them by name and saves them as one dictionary, which every loop from here to the end of the module relies on. When. Every model from this lesson on is built from nn.Linear layers, one here and several from Lesson 4.3.1.

import pandas as pd
import torch
from torch import nn

listings = pd.read_csv("listings.csv")
has_area = listings.dropna(subset=["area_sqm"])
area = torch.tensor(has_area["area_sqm"].to_numpy(), dtype=torch.float32)
price = torch.tensor(has_area["price"].to_numpy(), dtype=torch.float32)
z = (area - area.mean()) / area.std(correction=0)
X = z.unsqueeze(1)

torch.manual_seed(42)
model = nn.Linear(1, 1)
print(X.shape, X.dtype)
print(model)
print(model.weight)
print(model.bias)

torch.Size([2875, 1]) torch.float32
Linear(in_features=1, out_features=1, bias=True)
Parameter containing:
tensor([[0.7645]], requires_grad=True)
Parameter containing:
tensor([0.8300], requires_grad=True)

X is z with one dimension added by Lesson 4.1.2's unsqueeze, 2,875 rows of one feature, because the layer multiplies an (n, in_features) matrix by its transposed weight. torch.manual_seed(42) sits immediately before the layer because nn.Linear draws both values at creation, its initialization (the random starting values a layer draws when it is created), each between -1 / sqrt(in_features) and 1 / sqrt(in_features), so between -1 and 1 for one input, and the draw is 0.7645 and 0.8300 on every run of this release on a processor. Both print as Parameter containing: over a tensor with requires_grad=True that nobody wrote: a parameter asks for a gradient the moment it is made. The weight is (1, 1), a matrix with one entry, and the bias (1,), because the class is written for a layer of out_features neurons on in_features inputs, and one neuron on one input is the case where both counts are 1.

print(type(model.weight))
print(model.weight.is_leaf, model.weight.requires_grad, model.weight.grad)
for name, p in model.named_parameters():
    print(name, p.shape)
print(model.state_dict())
print(sum(p.numel() for p in model.parameters()))

<class 'torch.nn.parameter.Parameter'>
True True None
weight torch.Size([1, 1])
bias torch.Size([1])
OrderedDict({'weight': tensor([[0.7645]]), 'bias': tensor([0.8300])})
2

nn.Parameter is a subclass of torch.Tensor, so everything Lesson 4.1.4 established about a leaf holds: is_leaf is True, requires_grad is True, and .grad is None until a backward(). named_parameters() yields each parameter with its name, parameters() yields them without, which is the form an optimizer takes, and state_dict() returns the model's state dict (an ordered dictionary of every parameter's name and values, without their graphs), the form a model is saved in from Lesson 4.4.3. numel() counts a tensor's entries, so this model has 2 parameters, a number Lesson 4.3.1 makes much larger.

out = model(X)
print(out.shape, out.grad_fn.name())  # torch.Size([2875, 1]) AddmmBackward0
by_hand = X @ model.weight.T + model.bias
print(torch.allclose(out, by_hand))   # True
# True, 4.2.1's w * z + b with the layer's two values in place of w and b
print(torch.allclose(out.squeeze(1), model.weight[0, 0] * z + model.bias[0]))

model(X) runs the layer's forward and records it as one node, AddmmBackward0 (addmm, a matrix product and an addition in one operation), where Lesson 4.2.1's w * z + b recorded a MulBackward0 and an AddBackward0. The output has the shape of the target it will be scored against, (2875, 1), and its numbers are 4.2.1's w * z + b to float32's digits, with the layer's weight and bias in place of w and b. Two inputs are refused, each with a message worth reading once:

try:
    model(z)
except RuntimeError as err:
    print(err)  # mat1 and mat2 shapes cannot be multiplied (1x2875 and 1x1)
try:
    model(X.double())
except RuntimeError as err:
    print(err)  # mat1 and mat2 must have the same dtype, but got Double and Float

The one-dimensional z is read as a single row of 2,875 features, which no (1, 1) weight can multiply, so the message names the two shapes as matrices, 1x2875 and 1x1, with the inner dimensions, 2,875 and 1, the ones that do not match. A float64 input meets float32 parameters, and a matrix product refuses two dtypes where Lesson 4.1.2's elementwise arithmetic would have promoted, so the data is built float32 as it is above, Lesson 4.1.1's dtype for anything a model computes with, and a check against NumPy converts a copy with .double(), as Lesson 4.1.4 converted a gradient.

DON'T:

• DON'T hand the layer a one-dimensional tensor. model(z) is refused with mat1 and mat2 shapes cannot be multiplied (1x2875 and 1x1), and z.unsqueeze(1) is the (2875, 1) input it takes.
• DON'T build the input in float64 to match a NumPy check. model(X.double()) is refused with mat1 and mat2 must have the same dtype, but got Double and Float. The model's data stays float32, and the check converts its own copy.
• DON'T seed after the layer. nn.Linear(1, 1) draws at creation, so a torch.manual_seed(42) placed under it changes nothing about it, and a second layer created after the first draws the next values in the stream, -0.2343 and 0.9186, not 0.7645 and 0.8300 again.

nn.MSELoss Is the Loss, and the Shape It Wants

What. nn.MSELoss() is a module with no parameters whose call computes ((target - out) ** 2).mean(), Lesson 3.1.3's loss and Lesson 4.2.1's loss line, and it wants a target of the output's shape, (n, 1). Why. The loss as an object is what makes a loop generic: Lesson 4.3.2 swaps it for nn.CrossEntropyLoss and nn.BCEWithLogitsLoss, each with a target shape and dtype of its own, and the loop around it stays. The shape it wants is the one rule of this lesson that fails quietly, so it is shown failing, on a model that has already been fitted. When. Every regression loop in this module, on price.unsqueeze(1).

loss_fn = nn.MSELoss()
y = price.unsqueeze(1)
loss = loss_fn(out, y)
print(loss)                     # tensor(1.8286e+11, grad_fn=<MseLossBackward0>)
print(f"{loss.item():.4e}")     # 1.8286e+11
hand = ((y - out) ** 2).mean()
print(torch.equal(loss, hand))  # True

y is the price with the output's second dimension, the module's convention from here on: the target has the output's shape. loss_fn(out, y) is the loss Lesson 4.2.1 computed at a start near zero, 1.8286e+11, with a node of its own, MseLossBackward0, and torch.equal says it is the same float32 number as the hand formula, bit for bit. Handed the (2875,) price instead of y, the same call runs, warns and returns a number:

wrong = loss_fn(out, price)
print(f"{loss.item():.4e} {wrong.item():.4e}")
print((out - price).shape)
silent = ((price - out) ** 2).mean()
print(torch.equal(wrong, silent))

/home/goku/dev/bootcamp/.venv/lib/python3.12/site-packages/torch/nn/modules/loss.py:630: UserWarning: Using a target size (torch.Size([2875])) that is different to the input size (torch.Size([2875, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
  return F.mse_loss(input, target, reduction=self.reduction)
1.8286e+11 1.8286e+11
torch.Size([2875, 2875])
True

The warning comes from PyTorch's own loss.py, the line inside MSELoss.forward that calls the function, not from a line of yours, and the program goes on. Broadcasting, Lesson 4.1.2's table, stretches (2875, 1) against (2875,) to (2875, 2875): every prediction against every price, 8,265,625 squared differences averaged where the loss averages 2,875. At this start the two numbers agree to four digits because the predictions sit between -0.47 and 3.94, so every pair's difference is a price give or take four units, and the mean of every price squared is the same whether each price is counted once or 2,875 times. The hand-written ((price - out) ** 2).mean() broadcasts the same way to the same number and warns about nothing, because the warning is MSELoss's and not broadcasting's. The two numbers split once the model fits, so the model is fitted here, by Lesson 4.2.1's five steps on the module's two parameters, with the step and the clearing written by hand for the last time:

lr = 0.1
for epoch in range(1, 101):
    out = model(X)
    loss = loss_fn(out, y)
    loss.backward()
    with torch.no_grad():
        for p in model.parameters():
            p -= lr * p.grad
    for p in model.parameters():
        p.grad = None
    if epoch in (1, 10, 20, 30, 100):
        print(epoch, f"{loss.item():.4e}")
print(f"{model.weight.item():,.2f} {model.bias.item():,.2f}")

1 1.8286e+11
10 2.0209e+10
20 1.7260e+10
30 1.7226e+10
100 1.7226e+10
144,879.59 380,325.84

for p in model.parameters() is the whole reason a module exists: the loop steps and clears every parameter without naming one, and the losses are Lesson 4.2.1's, 2.0209e+10 at epoch 10, 1.7226e+10 from epoch 30, with the parameters at 144,879.59 and 380,325.84, that lesson's 144,880 and 380,326 to the unit. Now the two numbers part:

with torch.no_grad():
    out = model(X)
    right = loss_fn(out, y)
    wrong = loss_fn(out, price)
# 1.7226e+10 5.9206e+10
print(f"{right.item():.4e} {wrong.item():.4e}")
# 133,636 970,512
print(f"{out.min().item():,.0f} {out.max().item():,.0f}")

No second warning prints: Python's default filter shows each distinct warning message once per source line, and this second call raises the same message from the same loss.py:630, where a mismatch of another size, (100,) against (100, 1), is a different message and prints again. So the second mismatch prints only its number, 5.9206e+10 on the parameters whose loss is 1.7226e+10, 3.4 times the right one. The fitted predictions run from 133,636 to 970,512, and the broadcast scores the right prediction for a large listing against the price of every small one, so the better the predictions spread with the area, the further the pairwise mean sits from the loss. A loop that minimizes that number is minimizing the wrong thing, and Q4 of the Theory Assessment runs it.

DON'T:

• DON'T read the warning as noise because the loss still prints and falls. It does not stop the program, the loop runs to its end, and the number it drives down is not the loss: 5.9206e+10 against 1.7226e+10 on the same parameters.
• DON'T write the loss by hand to be rid of the warning. ((price - out) ** 2).mean() on the same two shapes computes the same wrong number and warns about nothing. The fix is the target's shape, price.unsqueeze(1), once, where the target is built.

optim.SGD Is the Step and the Zeroing

What. An optimizer (an object that holds a model's parameters and the rule that updates them from their gradients) is torch.optim.SGD(model.parameters(), lr=0.1) here. optimizer.step() runs p.add_(p.grad, alpha=-lr) on every parameter it holds, Lesson 4.2.1's step inside no_grad, and optimizer.zero_grad() sets every .grad to None, the fifth step. Why. The two steps the optimizer takes over are the two Lesson 4.1.4 showed failing quietly, the step on the record and the gradient never cleared. Inside the optimizer the step cannot be written on the record, the zeroing is one call that names no parameter, and a loop with ten layers of parameters is the same two lines as a loop with one neuron. When. Every loop from here on. SGD is the only optimizer of this section, and Lesson 4.3.4 puts Adam beside it.

import pandas as pd
import torch
from torch import nn

listings = pd.read_csv("listings.csv")
has_area = listings.dropna(subset=["area_sqm"])
area = torch.tensor(has_area["area_sqm"].to_numpy(), dtype=torch.float32)
price = torch.tensor(has_area["price"].to_numpy(), dtype=torch.float32)
z = (area - area.mean()) / area.std(correction=0)
X, y = z.unsqueeze(1), price.unsqueeze(1)

torch.manual_seed(42)
model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
print(optimizer)

SGD (
Parameter Group 0
    dampening: 0
    differentiable: False
    foreach: None
    fused: None
    lr: 0.1
    maximize: False
    momentum: 0
    nesterov: False
    weight_decay: 0
)

SGD is stochastic gradient descent (gradient descent where each step's gradient comes from a random part of the rows), and with every row in every step, as here and everywhere in this section, it is Lesson 4.2.1's gradient descent exactly, the random part arriving with Lesson 4.3.3's mini-batches. The optimizer took model.parameters() and the learning rate, and it prints one parameter group (a set of parameters that share one set of settings, the list optimizer.param_groups holds, with one entry here) with nine settings. lr is the one this section sets. momentum is Lesson 4.3.4's, with dampening and nesterov, which act only once momentum is not 0. weight_decay is Lesson 4.4.1's. maximize, foreach, fused and differentiable are switches no lesson of this section touches. One step, compared with the step Lesson 4.2.1 wrote:

before = model.weight.detach().clone()
loss = loss_fn(model(X), y)
loss.backward()
grad = model.weight.grad.clone()
print(f"{grad.item():.4e}")                 # -2.8976e+05
optimizer.step()
print(f"{model.weight.item():,.2f}")        # 28,976.56
fused = before.add(grad, alpha=-0.1)
plain = before - 0.1 * grad
print(torch.equal(model.weight, fused))     # True
print(torch.allclose(model.weight, plain))  # True
print(f"{model.weight.grad.item():.4e}")    # -2.8976e+05
optimizer.zero_grad()
print(model.weight.grad, model.bias.grad)   # None None

before is a detached copy of the weight, Lesson 4.1.4's .detach() and a clone() so the copy keeps its own memory and the step does not move it. The gradient at this start is -2.8976e+05, and one step at 0.1 moves the weight to 28,976.56. torch.equal says the parameter and before.add(grad, alpha=-0.1) are the same float32 number, because that is the line step() runs, and torch.allclose says the hand step, before - 0.1 * grad, lands within float32's tolerance of it. .grad is still there after the step, step() reads it and never clears it, and zero_grad() sets both gradients to None, the clearing Lesson 4.1.4 wrote as w.grad = None. Why the hand step is compared through allclose, on a thousand random elements instead of two:

torch.manual_seed(42)
p = torch.randn(1000)
g = torch.randn(1000)
fused = p.add(g, alpha=-0.1)
plain = p - 0.1 * g
print(torch.equal(fused, plain))          # False
print((fused != plain).sum())             # tensor(93)
print(torch.allclose(fused, plain))       # True
exact = p.double() - torch.tensor(0.1).double() * g.double()
print(torch.equal(fused, exact.float()))  # True
print(torch.equal(plain, exact.float()))  # False

add with alpha multiplies and adds in one rounding, and lands on the float32 nearest the exact result, exact.float(), on all 1,000 elements. p - 0.1 * g rounds the product, then the difference, and lands one or more float32 steps away on 93 of them. The two are the same arithmetic and different float32 numbers, which is the whole case for allclose on a parameter, Lesson 4.1.4's rule for a gradient extended to the step. One more fact about step() is worth one run:

torch.manual_seed(42)
model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
optimizer.step()
# 0.7645 0.8300
print(f"{model.weight.item():.4f} {model.bias.item():.4f}")

step() before any backward() finds every .grad at None, skips the parameter and says nothing: the weight and bias are the seeded draw, untouched. A loop fails this way at every epoch whenever a zero_grad() falls between each backward() and the next step(), as in zero_grad(), step(), backward(), while step(), zero_grad(), backward() moves from its second epoch on, each step reading the gradient the epoch before left.

DON'T:

• DON'T call step() before backward(), or zero_grad() between them. Every order that clears between a backward() and the next step() leaves .grad at None when the step reads it, nothing moves, nothing warns, and the loss prints 1.8286e+11 at every epoch: thirty epochs of backward(), zero_grad(), step() from this start ended on the seeded 0.7645 and 0.8300.
• DON'T drop zero_grad() because the optimizer handles the gradients. It handles them when called, and the module loop without the call prints 5.1426e+10 at epoch 5, 1.8034e+11 at epoch 15 and 1.1249e+11 at epoch 30 from this start, Lesson 4.2.1's bounce.
• DON'T test a parameter with == after a step. fused and plain differed on 93 of 1,000 elements while every printed digit agreed, and allclose is the comparison.

The Same Numbers, Proved

What. Two loops from one seeded start, Lesson 4.2.1's hand loop restated as that lesson wrote it and the module loop of the three classes, their losses printed side by side and their parameters compared with allclose. Why. A refactor is honest when the numbers do not move. Without the proof, optimizer.step() is a line you trust, and with it, it is a line you have watched do what w -= lr * w.grad did, 100 epochs running. When. Once, here, and again whenever a line you wrote is replaced by a library call that claims to do the same: the test is the same numbers from the same start.

import pandas as pd
import torch
from torch import nn

listings = pd.read_csv("listings.csv")
has_area = listings.dropna(subset=["area_sqm"])
area = torch.tensor(has_area["area_sqm"].to_numpy(), dtype=torch.float32)
price = torch.tensor(has_area["price"].to_numpy(), dtype=torch.float32)
z = (area - area.mean()) / area.std(correction=0)
X, y = z.unsqueeze(1), price.unsqueeze(1)
lr = 0.1

torch.manual_seed(42)
model = nn.Linear(1, 1)
w = torch.tensor(model.weight.item(), requires_grad=True)
b = torch.tensor(model.bias.item(), requires_grad=True)
print(w, b)  # tensor(0.7645, requires_grad=True) tensor(0.8300, requires_grad=True)

hand = []
for epoch in range(1, 101):
    pred = w * z + b
    loss = ((price - pred) ** 2).mean()
    loss.backward()
    with torch.no_grad():
        w -= lr * w.grad
        b -= lr * b.grad
    w.grad = None
    b.grad = None
    hand.append(f"{loss.item():.4e}")

The hand loop is Lesson 4.2.1's, step for step, from the module's seeded start instead of zero: w and b are made from model.weight.item() and model.bias.item(), the same float32 values, so both loops begin at 0.7645 and 0.8300. Each epoch's loss is kept as its printed string, .4e, the precision the proof compares at. The module loop follows on the same model, which no step has touched yet:

loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
module = []
for epoch in range(1, 101):
    out = model(X)
    loss = loss_fn(out, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    module.append(f"{loss.item():.4e}")

for epoch in (1, 2, 3, 10, 20, 30, 70, 100):
    print(epoch, hand[epoch - 1], module[epoch - 1])
print(sum(h == m for h, m in zip(hand, module)))

1 1.8286e+11 1.8286e+11
2 1.2323e+11 1.2323e+11
3 8.5070e+10 8.5070e+10
10 2.0209e+10 2.0209e+10
20 1.7260e+10 1.7260e+10
30 1.7226e+10 1.7226e+10
70 1.7226e+10 1.7226e+10
100 1.7226e+10 1.7226e+10
100

The two strings are the same at all 100 epochs and the count prints 100. They match Lesson 4.2.1's loop from zero at 98 epochs and sit one lower in the last digit at two, 8.5070e+10 at epoch 3 against its 8.5071e+10 and 6.0646e+10 at epoch 4 against its 6.0647e+10, the one trace of the seeded start in the hundred strings. The module loop is five lines and 4.2.1's five steps in order: model(X) predicts, loss_fn(out, y) measures, loss.backward() walks the graph, optimizer.step() steps inside no_grad, and optimizer.zero_grad() clears. From here on that is how you read any PyTorch loop, five steps, whatever the model, the loss and the optimizer are called, and a loop that moves the zeroing to the top of its body is the same loop, the 100 strings and the parameters unchanged from a run. The parameters:

print(f"{w.item():,.2f} {b.item():,.2f}")  # 144,879.59 380,325.84
# 144,879.59 380,325.84
print(f"{model.weight.item():,.2f} {model.bias.item():,.2f}")
# True True
print(torch.allclose(model.weight, w), torch.allclose(model.bias, b))
per_sqm = model.weight.item() / area.std(correction=0).item()
print(round(per_sqm, 2))                   # 3288.31

Both loops end at 144,879.59 and 380,325.84, least squares' 144,879.63 and 380,325.91 within float32's reach, Lesson 4.2.1's 144,880 and 380,326 to the unit, and allclose holds for both parameters, each broadcast against the hand scalar. The weight reads back as 4.2.1 read it, 3,288.31 per square meter. The classes computed what you wrote and nothing more: the module loop is the hand loop with the leaves in a layer and the step in an optimizer.

DON'T:

• DON'T read the 100 as ==. The strings agree at four significant digits and allclose holds on the parameters, and that is the proof: the last bits of a float32 parameter belong to the step's rounding, not to the fit.
• DON'T prove a refactor on one epoch. The first loss agrees whenever both loops read the same parameters, and the step is what the refactor replaced, so the proof runs the loop out and compares where it ends.

Your Own nn.Module

What. A module of your own is a class that subclasses nn.Module, calls super().__init__() first, assigns its layers as attributes in __init__, and defines forward(self, x). Neuron holds one nn.Linear(1, 1), and model(x) runs its forward. Why. nn.Linear alone fits this section's model, and Lesson 4.3.1's network is this class with more layers and a function between them, so the shape is learned on one layer, where every number can be checked against the subtitle above. When. Whenever a model is more than one layer, which is every model from Section 4.3 on, and here, once, so that Lesson 4.3.1's class is a change of two lines.

import pandas as pd
import torch
from torch import nn

listings = pd.read_csv("listings.csv")
has_area = listings.dropna(subset=["area_sqm"])
area = torch.tensor(has_area["area_sqm"].to_numpy(), dtype=torch.float32)
price = torch.tensor(has_area["price"].to_numpy(), dtype=torch.float32)
z = (area - area.mean()) / area.std(correction=0)
X, y = z.unsqueeze(1), price.unsqueeze(1)


class Neuron(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)


torch.manual_seed(42)
model = Neuron()
print(model)
for name, p in model.named_parameters():
    print(name, p.shape)
print(f"{model.linear.weight.item():.4f} {model.linear.bias.item():.4f}")
print(model.training)

Neuron(
  (linear): Linear(in_features=1, out_features=1, bias=True)
)
linear.weight torch.Size([1, 1])
linear.bias torch.Size([1])
0.7645 0.8300
True

super().__init__() is what makes the class a module: it sets up the bookkeeping that notices every nn.Module or nn.Parameter assigned as an attribute and registers it, so print(model) lists the layer under its attribute name and the parameters arrive with that name in front, linear.weight and linear.bias. The seed immediately before Neuron() gives the layer the same draw as before, 0.7645 and 0.8300, because the only random draw in the constructor is the layer's. training is a flag every module carries, True at creation, which nothing in this model reads. The class is Section 1.8's class in Lesson 3.1.1's estimator shape, which Lesson 3.1.2's Memorizer filled with fit and predict, with forward as the predict and the fit left to the loop:

loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
for epoch in range(1, 101):
    out = model(X)
    loss = loss_fn(out, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    if epoch in (1, 10, 100):
        print(epoch, f"{loss.item():.4e}")
print(f"{model.linear.weight.item():,.2f} {model.linear.bias.item():,.2f}")
print(list(model.state_dict().keys()))
model.eval()
print(model.training, model.linear.training)
model.train()
print(model.training)

1 1.8286e+11
10 2.0209e+10
100 1.7226e+10
144,879.59 380,325.84
['linear.weight', 'linear.bias']
False False
True

Same five lines, same losses to the digit, same parameters, because the only computation in forward is the layer's. model.parameters() reaches through to linear.weight and linear.bias, which is why the optimizer takes a module of any depth with one call, and the state dict carries the same two names. model.eval() and model.train() flip training on the module and on every module inside it: nothing here reads the flag, and Lesson 4.2.3 says when to flip it. Two refusals, from a class missing one of its two required lines:

class NoSuper(nn.Module):
    def __init__(self):
        self.linear = nn.Linear(1, 1)


class NoForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)


try:
    NoSuper()
except AttributeError as err:
    print(err)
try:
    NoForward()(X)
except NotImplementedError as err:
    print(err)

cannot assign module before Module.__init__() call
Module [NoForward] is missing the required "forward" function

Without super().__init__(), the first attribute assignment finds no bookkeeping and raises before the constructor ends. Without forward, the class builds and the call raises, naming the class, because nn.Module.__call__ is what runs forward, with the module's hooks (functions a module can register to run before or after forward) around it. That is why a model is called as model(x) and never as model.forward(x): the numbers are the same here, and the hooks are what the second form skips.

DON'T:

• DON'T assign a layer before super().__init__(). self.linear = nn.Linear(1, 1) on the first line of __init__ raises cannot assign module before Module.__init__() call, and the call goes first in every module you write.
• DON'T store a parameter as a plain tensor attribute. self.w = torch.zeros(1, 1) in __init__ is not in parameters(), so the optimizer never sees it, while nn.Parameter(torch.zeros(1, 1)) and any layer are registered as they are assigned.
• DON'T call model.forward(x). It returns the same tensor as model(x) on this model and skips __call__, where the hooks live, and every loop in this module calls the model.

Lab Exercises

Lab Exercise 1: A Layer on the Room Count

In rooms_linear.py, load listings.csv and use all 3,000 listings, since no listing lacks a room count. Build float32 tensors of rooms and price, standardize the rooms with their mean and std(correction=0), and build X of shape (3000, 1) and y of shape (3000, 1). Seed with 42, create nn.Linear(1, 1), and print the model, its weight and bias to four decimals, then the output's shape beside the nn.MSELoss loss to four significant digits, and whether that loss equals ((y - out) ** 2).mean() under torch.equal. Catch and print the two refusals, the one-dimensional input and the float64 input.

Then run Lesson 4.2.1's five steps on model.parameters() by hand, with the step inside torch.no_grad() and the gradients set to None, for 100 epochs at a learning rate of 0.1, printing the loss at epochs 1, 20, 40 and 100 and the fitted weight and bias to two decimals. Check them against np.linalg.lstsq on the standardized rooms and a column of ones with np.allclose at rtol=1e-4, and read the weight back per room. End by computing, under no_grad, the loss against y and against the (3000,) price, printed side by side.

Lab Exercise 2: A Neuron Class on the Distance

In distance_neuron.py, on the 2,875 listings with an area, build float32 tensors of distance_km and price, standardize the distance, and build X and y of shape (2875, 1). Write Neuron, a subclass of nn.Module holding one nn.Linear(1, 1), seed with 42, create it, build torch.optim.SGD at 0.1, and print the names of its parameters. Call optimizer.step() once before any backward() and print the weight to four decimals. Then take one real step: compute the loss, call backward(), keep a detached clone of the weight and of its gradient, call step(), and print whether the weight equals the clone's add(grad, alpha=-0.1) under torch.equal and whether it is allclose to the clone minus 0.1 * grad. Print the gradient still in .grad to four significant digits, call zero_grad(), and print .grad.

Then create a fresh seeded Neuron and run the module loop for 100 epochs, printing the loss at epochs 1, 20, 40 and 100 and the fitted parameters to two decimals. Check them against np.linalg.lstsq with np.allclose at rtol=1e-4, and read the weight back per kilometer.

import numpy as np
import pandas as pd
import torch
from torch import nn

listings = pd.read_csv("listings.csv")
rooms = torch.tensor(listings["rooms"].to_numpy(), dtype=torch.float32)
price = torch.tensor(listings["price"].to_numpy(), dtype=torch.float32)
z = (rooms - rooms.mean()) / rooms.std(correction=0)
X, y = z.unsqueeze(1), price.unsqueeze(1)

torch.manual_seed(42)
model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
print(model)
print(f"{model.weight.item():.4f} {model.bias.item():.4f}")
out = model(X)
loss = loss_fn(out, y)
print(out.shape, f"{loss.item():.4e}")
print(torch.equal(loss, ((y - out) ** 2).mean()))
try:
    model(z)
except RuntimeError as err:
    print(err)
try:
    model(X.double())
except RuntimeError as err:
    print(err)

Linear(in_features=1, out_features=1, bias=True)
0.7645 0.8300
torch.Size([3000, 1]) 1.8311e+11
True
mat1 and mat2 shapes cannot be multiplied (1x3000 and 1x1)
mat1 and mat2 must have the same dtype, but got Double and Float

lr = 0.1
for epoch in range(1, 101):
    out = model(X)
    loss = loss_fn(out, y)
    loss.backward()
    with torch.no_grad():
        for p in model.parameters():
            p -= lr * p.grad
    for p in model.parameters():
        p.grad = None
    if epoch in (1, 20, 40, 100):
        print(epoch, f"{loss.item():.4e}")
print(f"{model.weight.item():,.2f} {model.bias.item():,.2f}")

A = np.column_stack([z.numpy().astype(np.float64), np.ones(len(z))])
w_ls, b_ls = np.linalg.lstsq(A, price.numpy().astype(np.float64), rcond=None)[0]
print(f"{w_ls:,.2f} {b_ls:,.2f}")
fitted = [model.weight.item(), model.bias.item()]
print(np.allclose(fitted, [w_ls, b_ls], rtol=1e-4))
print(round(fitted[0] / rooms.std(correction=0).item(), 2))
with torch.no_grad():
    out = model(X)
    print(f"{loss_fn(out, y).item():.4e} {loss_fn(out, price).item():.4e}")

1 1.8311e+11
20 2.0501e+10
40 2.0467e+10
100 2.0467e+10
132,899.22 380,762.94
132,899.24 380,763.00
True
77791.18
/home/goku/dev/bootcamp/.venv/lib/python3.12/site-packages/torch/nn/modules/loss.py:630: UserWarning: Using a target size (torch.Size([3000])) that is different to the input size (torch.Size([3000, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
  return F.mse_loss(input, target, reduction=self.reduction)
2.0467e+10 5.5792e+10

import numpy as np
import pandas as pd
import torch
from torch import nn

listings = pd.read_csv("listings.csv")
has_area = listings.dropna(subset=["area_sqm"])
dist = torch.tensor(has_area["distance_km"].to_numpy(), dtype=torch.float32)
price = torch.tensor(has_area["price"].to_numpy(), dtype=torch.float32)
z = (dist - dist.mean()) / dist.std(correction=0)
X, y = z.unsqueeze(1), price.unsqueeze(1)


class Neuron(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)


torch.manual_seed(42)
model = Neuron()
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
print([name for name, p in model.named_parameters()])
optimizer.step()
print(f"{model.linear.weight.item():.4f}")

before = model.linear.weight.detach().clone()
loss = loss_fn(model(X), y)
loss.backward()
grad = model.linear.weight.grad.clone()
optimizer.step()
print(torch.equal(model.linear.weight, before.add(grad, alpha=-0.1)))
print(torch.allclose(model.linear.weight, before - 0.1 * grad))
print(f"{model.linear.weight.grad.item():.4e}")
optimizer.zero_grad()
print(model.linear.weight.grad)

['linear.weight', 'linear.bias']
0.7645
True
True
-8.1088e+03
None

torch.manual_seed(42)
model = Neuron()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
for epoch in range(1, 101):
    out = model(X)
    loss = loss_fn(out, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    if epoch in (1, 20, 40, 100):
        print(epoch, f"{loss.item():.4e}")
print(f"{model.linear.weight.item():,.2f} {model.linear.bias.item():,.2f}")

A = np.column_stack([z.numpy().astype(np.float64), np.ones(len(z))])
w_ls, b_ls = np.linalg.lstsq(A, price.numpy().astype(np.float64), rcond=None)[0]
print(f"{w_ls:,.2f} {b_ls:,.2f}")
fitted = [model.linear.weight.item(), model.linear.bias.item()]
print(np.allclose(fitted, [w_ls, b_ls], rtol=1e-4))
print(round(fitted[0] / dist.std(correction=0).item(), 2))

1 1.8286e+11
20 3.8229e+10
40 3.8199e+10
100 3.8199e+10
4,055.12 380,325.84
4,055.12 380,325.91
True
1331.63

Theory Assessment

Q1: What a Layer Holds

nn.Linear(1, 1) built after torch.manual_seed(42) printed a weight of 0.7645 and a bias of 0.8300. Say what the two tensors are, with their shapes and why those shapes, why both already require a gradient and are leaves, and what model(X) computes for the (2875, 1) input. Then say why model(z) on the (2875,) tensor is refused and what the 1x2875 in its message is.

Hint: The class is written for out_features neurons on in_features inputs, and one neuron on one input is the case where both are 1. For the refusal, ask what shape a one-dimensional tensor is read as when it meets a matrix product.

Q2: The Number the Warning Guards

At the seeded start, loss_fn(out, y) and loss_fn(out, price) both printed 1.8286e+11, and a teammate concludes the warning is noise. Say what the second call computes and how many squared differences it averages, why the two agree to four digits at the start and what they printed on the fitted model, and what the hand-written ((price - out) ** 2).mean() does with the same two shapes.

Hint: Write the two shapes, (2875, 1) and (2875,), and apply Lesson 4.1.2's broadcasting rule. For the start, ask how far apart the predictions were, and for the fitted model, how far apart they are now.

Q3: What the Optimizer Took Over

Of Lesson 4.2.1's five steps, say which optimizer.step() takes over, which optimizer.zero_grad() takes over, and which three stay written in the loop. Say what step() does to a parameter whose .grad is None, what zero_grad() leaves in .grad, and why a parameter after step() is compared with the hand step through allclose and not through ==.

Hint: Count the lines of the module loop. For the comparison, ask how many roundings p.add_(g, alpha=-lr) performs and how many p - lr * g performs.

Q4: The Loop That Fit the Wrong Loss

Read this script without running it, then say what it prints and what the model has learned.

import pandas as pd
import torch
from torch import nn

listings = pd.read_csv("listings.csv")
has_area = listings.dropna(subset=["area_sqm"])
area = torch.tensor(has_area["area_sqm"].to_numpy(), dtype=torch.float32)
price = torch.tensor(has_area["price"].to_numpy(), dtype=torch.float32)
z = (area - area.mean()) / area.std(correction=0)
X = z.unsqueeze(1)

torch.manual_seed(42)
model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
for epoch in range(1, 101):
    out = model(X)
    loss = loss_fn(out, price)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    if epoch in (1, 10, 20, 30, 100):
        print(epoch, f"{loss.item():.4e}")
print(f"{model.weight.item():,.2f} {model.bias.item():,.2f}")

A teammate says the loop settled, the loss stopped moving by epoch 30, so the model fits, a little worse than Lesson 4.2.1's 1.7226e+10 because nn.Linear starts from a random draw instead of zero. Say what the loss settles at, what the two parameters print, and where each part of that reading goes wrong.

Hint: The target is price, not y. Ask which number the loop is driving down, and which single prediction makes every pair's squared difference as small as it can be at once.

/home/goku/dev/bootcamp/.venv/lib/python3.12/site-packages/torch/nn/modules/loss.py:630: UserWarning: Using a target size (torch.Size([2875])) that is different to the input size (torch.Size([2875, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
  return F.mse_loss(input, target, reduction=self.reduction)
1 1.8286e+11
10 4.0821e+10
20 3.8246e+10
30 3.8216e+10
100 3.8216e+10
0.00 380,325.84

Assignment

Assignment: Two Features, One Layer

Your pricing model has used the area alone since Lesson 3.1.3, and the first thing a teammate asks of the refactored loop is whether it takes a second feature without a new line in the loop. Write two_features.py on the 2,875 listings with an area. Standardize area_sqm and distance_km each by its own mean and std(correction=0), and build the input with torch.stack([z_area, z_dist], dim=1), which this lesson did not use: it places two (n,) tensors side by side as the columns of one (n, 2) tensor. nn.Linear(2, 1), seeded with 42, is the layer with two inputs, a weight of shape (1, 2) beside the same (1,) bias, and its model(X) is still X @ weight.T + bias, Lesson 2.3.2's batch prediction with two weights. Print the input's shape and the seeded weight and bias.

Before any step, compute the loss on y, call backward(), and print model.weight.grad's shape and its two entries beside model.bias.grad, to four significant digits. Check them with torch.allclose against Lesson 4.1.3's partial derivatives computed by hand on the detached residual, mean(-2 * z_k * (price - pred)) for each feature and mean(-2 * (price - pred)) for the bias, stacked into the gradient's own shapes.

Train with torch.optim.SGD at 0.1 for 100 epochs, printing the loss at epochs 1, 25, 50 and 100 and the three parameters to two decimals, checked against np.linalg.lstsq on the two standardized columns and a column of ones with np.allclose at rtol=1e-4. Read the weights back in the units of the columns, w / sd for each, and the intercept as b - w_area * mean_area / sd_area - w_dist * mean_dist / sd_dist, and compare all three with LinearRegression fitted on the two raw columns, Lesson 3.1.1's model, through np.allclose at the same tolerance. Print the first three predictions of both models in whole units, and compare the distance weight with Lab Exercise 2's, where the distance was the only feature.

Hints:

• torch.stack([z_area, z_dist], dim=1) gives (2875, 2), and model.weight.grad has the weight's shape, (1, 2), so the two hand derivatives are stacked and given a first dimension, torch.stack([g_area, g_dist]).unsqueeze(0), before allclose, and the bias derivative takes unsqueeze(0) to match (1,).
• The check's backward() leaves a gradient in .grad, and the loop's first backward() would add to it, so optimizer.zero_grad() goes before the loop. LinearRegression reports coef_ per raw unit and intercept_, so the standardized intercept is converted before it is compared.

import numpy as np
import pandas as pd
import torch
from torch import nn
from sklearn.linear_model import LinearRegression

listings = pd.read_csv("listings.csv")
has_area = listings.dropna(subset=["area_sqm"])
area = torch.tensor(has_area["area_sqm"].to_numpy(), dtype=torch.float32)
dist = torch.tensor(has_area["distance_km"].to_numpy(), dtype=torch.float32)
price = torch.tensor(has_area["price"].to_numpy(), dtype=torch.float32)
z_area = (area - area.mean()) / area.std(correction=0)
z_dist = (dist - dist.mean()) / dist.std(correction=0)
X = torch.stack([z_area, z_dist], dim=1)
y = price.unsqueeze(1)

torch.manual_seed(42)
model = nn.Linear(2, 1)
print(X.shape)
print(model.weight)
print(model.bias)
loss_fn = nn.MSELoss()
out = model(X)
loss = loss_fn(out, y)
loss.backward()
g = model.weight.grad
print(g.shape, f"{g[0, 0].item():.4e} {g[0, 1].item():.4e} {model.bias.grad.item():.4e}")

# 4.1.3's partial derivatives, one per feature, on the residual with no graph
residual = (price - out.squeeze(1)).detach()
g_area = (-2 * z_area * residual).mean()
g_dist = (-2 * z_dist * residual).mean()
hand_w = torch.stack([g_area, g_dist]).unsqueeze(0)
hand_b = (-2 * residual).mean().unsqueeze(0)
print(f"{hand_w[0, 0].item():.4e} {hand_w[0, 1].item():.4e} {hand_b.item():.4e}")
print(torch.allclose(g, hand_w), torch.allclose(model.bias.grad, hand_b))

torch.Size([2875, 2])
Parameter containing:
tensor([[0.5406, 0.5869]], requires_grad=True)
Parameter containing:
tensor([-0.1657], requires_grad=True)
torch.Size([1, 2]) -2.8976e+05 -8.1089e+03 -7.6065e+05
-2.8976e+05 -8.1089e+03 -7.6065e+05
True True

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
optimizer.zero_grad()
for epoch in range(1, 101):
    out = model(X)
    loss = loss_fn(out, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    if epoch in (1, 25, 50, 100):
        print(epoch, f"{loss.item():.4e}")
w_area, w_dist = model.weight[0, 0].item(), model.weight[0, 1].item()
b = model.bias.item()
print(f"{w_area:,.2f} {w_dist:,.2f} {b:,.2f}")

A = np.column_stack([z_area.numpy(), z_dist.numpy(), np.ones(len(z_area))])
ls = np.linalg.lstsq(A.astype(np.float64), price.numpy().astype(np.float64), rcond=None)[0]
print(f"{ls[0]:,.2f} {ls[1]:,.2f} {ls[2]:,.2f}")
print(np.allclose([w_area, w_dist, b], ls, rtol=1e-4))

# back to the columns' units, 4.2.1's w / sd for each feature, and the raw intercept
per_sqm = w_area / area.std(correction=0).item()
per_km = w_dist / dist.std(correction=0).item()
intercept = b - per_sqm * area.mean().item() - per_km * dist.mean().item()
print(f"{per_sqm:,.2f} {per_km:,.2f} {intercept:,.2f}")
raw = has_area[["area_sqm", "distance_km"]].to_numpy(dtype=np.float64)
lin = LinearRegression().fit(raw, has_area["price"].to_numpy(dtype=np.float64))
print(f"{lin.coef_[0]:,.2f} {lin.coef_[1]:,.2f} {lin.intercept_:,.2f}")
print(np.allclose([per_sqm, per_km, intercept], [*lin.coef_, lin.intercept_], rtol=1e-4))
with torch.no_grad():
    print([round(v) for v in model(X[:3]).squeeze(1).tolist()])
print([round(v) for v in lin.predict(raw[:3]).tolist()])

1 1.8286e+11
25 1.5737e+10
50 1.5725e+10
100 1.5725e+10
156,366.23 -40,401.58 380,325.84
156,366.26 -40,401.60 380,325.92
True
3,549.02 -13,267.21 101,806.81
3,549.03 -13,267.21 101,806.85
True
[598005, 134429, 570318]
[598005, 134429, 570319]

Additional Resources

Optional, and nothing later depends on either. Worth opening if you want the full list of what a module and an optimizer can be asked to do beyond the calls this lesson used.

• Official docs - PyTorch 2.14, torch.nn.Module. Every method of the class every model subclasses, with its signature: parameters(), named_parameters(), state_dict() and load_state_dict(), train() and eval(), zero_grad(), apply(), to(), and the hooks, register_forward_hook() and its siblings.
• Official docs - PyTorch 2.14, torch.optim.SGD. The algorithm box that writes out six of the nine settings print(optimizer) showed, lr, weight_decay, momentum, dampening, nesterov and maximize, as the change each one makes to the update, the parameter list where foreach, fused and differentiable switch how the step is run and not what it computes, and the methods step(), zero_grad(), state_dict() and add_param_group().

Summary

The refactor turned out to be bookkeeping. A layer is a leaf with a name, a loss object is a loss line that checks shapes, and an optimizer is the step and the clearing with the parameters already in hand. Every number the module loop printed was the hand loop's number, at every epoch of a hundred and to two decimals of the parameters, because the three classes do what the hand loop did and nothing else. That is what a library is for, and it is also why the hand loop came first: you can read optimizer.step() as one line of arithmetic because you have seen the arithmetic.

What the classes do not do is think for you. The loss object warned and went on, and a loop that trusts a warning away fits the mean price of the city to every listing. The optimizer stepped when told and skipped a gradient that was not there without a word. The module registered what was assigned after its own __init__ and nothing before it. Each failure was a line in the wrong place or a tensor in the wrong shape, and each was caught here the same way, by a number compared with a number you could compute by hand.

Lesson 4.2.3 keeps this loop and changes what is scored: rows held out, a validation loss read beside the training loss, and the two curves drawn. Section 4.3 keeps the loop too and changes the model, the loss and the optimizer, one object at a time. The five steps do not change again.

Cheat Sheet

Part A - Reference Table

Part B - Compact Code Block

import torch
from torch import nn

torch.manual_seed(42)        # before the layer, never after
model = nn.Linear(1, 1)      # weight (1, 1), bias (1,)
loss_fn = nn.MSELoss()       # the loss as an object
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
X, y = z.unsqueeze(1), price.unsqueeze(1)

for epoch in range(1, 101):
    out = model(X)           # predict, (n, 1)
    loss = loss_fn(out, y)   # measure, y is (n, 1)
    loss.backward()          # walk the graph
    optimizer.step()         # p.add_(p.grad, alpha=-lr)
    optimizer.zero_grad()    # every .grad to None
print(f"{loss.item():.4e}")  # 1.7226e+10

class Neuron(nn.Module):
    def __init__(self):
        super().__init__()              # first, or the next line raises
        self.linear = nn.Linear(1, 1)   # registered as linear.weight, linear.bias

    def forward(self, x):
        return self.linear(x)           # run through model(x), never model.forward(x)

model = Neuron()
print(model)                            # Neuron( (linear): Linear(...) )
model.state_dict().keys()               # linear.weight, linear.bias
torch.allclose(model.linear.weight, w)  # the comparison, never ==