Getting started with PyTorch

In the first three lessons we wrote linear regression, gradient descent and backpropagation by hand. They work, but they're slow, and every new model means rewriting lots of code.

PyTorch is currently the most widely used deep learning framework. This lesson doesn't cover all its features; it does one thing: redo what the first three lessons did, with PyTorch, comparing every step with the hand-written version. You'll find it does exactly what you wrote by hand, just in a faster, more convenient way.

Installation (no graphics card needed; this lesson and the next module all run on a CPU):

uv add torch

Tensors and automatic differentiation

The most basic thing in PyTorch is the tensor: a multi-dimensional array, much like a numpy array. A single number is a 0-dimensional tensor, a list of numbers is 1-dimensional, a table is 2-dimensional.

Tensors have an ability numpy arrays lack: automatic differentiation. Create a tensor with requires_grad=True and PyTorch records every operation it takes part in, just like last lesson's Num:

a = torch.tensor(2.0, requires_grad=True)  # requires_grad：记下它参与的运算，以便求导
b = torch.tensor(-3.0, requires_grad=True)
c = torch.tensor(10.0, requires_grad=True)
f = (a * b + c) ** 2
f.backward()
print(f"  f = {f.item()}，df/da={a.grad.item()}, df/db={b.grad.item()}, df/dc={c.grad.item()}")

== 1. 同一个算式 f = (a × b + c)²，a=2, b=-3, c=10
  f = 16.0，df/da=-24.0, df/db=16.0, df/dc=8.0

Exactly the same results as last lesson's hand-written Num: -24, 16, 8. What f.backward() does is what last lesson's Num.backward did: starting from f, walk back along the recorded computation, pass the gradients back with the chain rule, and store them in each tensor's .grad. .item() converts a tensor holding a single number into an ordinary Python number.

Linear regression: four things

Rewriting Lesson 2's linear regression in PyTorch uses the four basic things PyTorch trains models with:

x = torch.tensor((area - area.mean()) / area.std(), dtype=torch.float32).unsqueeze(1)  # 标准化，形状 (50, 1)
y = torch.tensor(price, dtype=torch.float32).unsqueeze(1)

model = nn.Linear(1, 1)  # 就是 w × x + b
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.MSELoss()
for step in range(100):
    loss = loss_fn(model(x), y)
    optimizer.zero_grad()  # 清空上一步的梯度
    loss.backward()  # 反向传播，算出每个参数的梯度
    optimizer.step()  # 按梯度更新参数：p -= lr × p.grad

• The model: nn.Linear(1, 1) is w × x + b, 1 input and 1 output. It creates and manages the two parameters w and b itself, with requires_grad already set.
• The loss function: nn.MSELoss() is Lesson 1's mean squared error.
• The optimiser: torch.optim.SGD updates the parameters. optimizer.step() is Lesson 2's w -= lr × dw, except it does it for every parameter in model.parameters().
• The training loop: zero the gradients, backpropagate, update the parameters, three lines in exactly the same order as last lesson's hand-written version.

unsqueeze(1) changes the shape from (50,) to (50, 1), meaning "50 samples, 1 feature each". PyTorch layers all assume the first dimension is the number of samples; more on that below.

== 2. 线性回归 100 步后：损失 143.05，换算回原单位 w=1.2571 b=14.806

Exactly the same result as Lesson 2's hand-written gradient descent: loss 143.05, w = 1.2571, b = 14.806.

Your own network: subclass nn.Module

For a somewhat more complex model, write your own class that inherits from nn.Module. Here's last lesson's XOR network rewritten with it:

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(2, 4)
        self.output = nn.Linear(4, 1)

    def forward(self, x):
        return torch.tanh(self.output(torch.tanh(self.hidden(x))))

__init__ defines which layers there are, and forward says how data flows through them. Last lesson we wrote two classes, Neuron and Layer; now a single nn.Linear(2, 4) is "a layer of 4 neurons with 2 inputs each", combining the 4 neurons' computations into one matrix multiplication.

nn.Module automatically finds the parameters of every layer you define in __init__, and net.parameters() returns them. Count them:

== 3. 异或网络，共 17 个参数

17 parameters, the same as last lesson's hand-written network. Training:

X = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=torch.float32)
Y = torch.tensor([[-1], [1], [1], [-1]], dtype=torch.float32)

optimizer = torch.optim.SGD(net.parameters(), lr=0.1)
for epoch in range(1, 301):
    loss = ((net(X) - Y) ** 2).sum()  # 和第 3 课一样：4 个样本的平方误差之和
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Notice that net(X) computes all 4 samples at once. Last lesson we computed them one at a time in a for loop; now the 4 samples go into a tensor of shape (4, 2), and one matrix operation computes them all.

  第   1 轮  损失 4.5469
  第  10 轮  损失 3.7176
  第  50 轮  损失 0.5588
  第 100 轮  损失 0.0514
  第 200 轮  损失 0.0174
  第 300 轮  损失 0.0100
  训练后的预测： ['-0.953', '+0.964', '+0.944', '-0.943']

It learns XOR too. The loss doesn't change in quite the same way as last lesson's numbers, because the initial parameter values are random and PyTorch uses different random numbers from our hand-written version. But the trend is the same: slow at first, then a sudden drop, and finally close to 0.

When predicting after training, the code is wrapped in torch.no_grad():

with torch.no_grad():  # 只是预测，不需要记录求导信息
    print("  训练后的预测：", [f"{v:+.3f}" for v in net(X).squeeze(1).tolist()])

By default, PyTorch records the information needed for differentiation at every operation. When you're only predicting and don't need to train, turning that off saves memory and runs faster.

Shapes: the most error-prone part

When writing PyTorch code, most errors have to do with tensor shapes. One convention to remember: the first dimension is the number of samples in a batch.

batch = torch.randn(32, 2)  # 32 个样本，每个 2 个特征

== 4. 一次喂 32 个样本：输入形状 (32, 2) → 隐藏层 (32, 4) → 输出 (32, 1)

The input is 32 samples with 2 features each; after the hidden layer it's 32 samples with 4 numbers each; at the end it's 32 samples with 1 output each. The number of samples stays the same throughout; each layer changes only the last dimension.

Processing a batch of samples at once, rather than one at a time, is one of the key reasons deep learning is fast. Matrix operations are highly optimised on CPUs and especially on GPUs, and computing 32 samples at once takes hardly more time than computing 1.

When writing code, printing x.shape at key points is the most effective debugging technique.

Hand-written versus PyTorch

PyTorch does no magic. You've already written its most central parts yourself; now you're just switching to a faster, less laborious way of writing them. When you later hit strange training problems, like a loss that won't fall or gradients that turn into NaN, you'll know what's happening underneath, and so where to start looking.

Exercises

1. Switch the linear regression data to unstandardised area (use area directly). What learning rate avoids divergence? Does that agree with Lesson 2's conclusion?
2. Replace torch.optim.SGD with torch.optim.Adam, set the learning rate to 0.01, and retrain the XOR network. Does the loss fall faster or slower?
3. Deliberately cause a shape error: change the XOR input X to shape (4, 3) (3 features per sample), run it, and read the error message to see how it tells you the shapes don't match.

Self-check

1. What do optimizer.zero_grad(), loss.backward() and optimizer.step() each do?

zero_grad resets every parameter's gradient to zero so it doesn't accumulate with the previous step's; backward backpropagates from the loss, computing each parameter's gradient and storing it in .grad; step updates each parameter according to its gradient, which for the simplest SGD is p = p - lr × p.grad.

2. Why does PyTorch assume a tensor's first dimension is the number of samples?

Deep learning always processes a batch of samples at once, putting them in one tensor and computing them together with matrix operations, which is much faster than one at a time. With the first dimension fixed as the number of samples, each layer only has to handle the final feature dimension, and the number of samples stays the same throughout the network.

3. Why use torch.no_grad() when predicting?

By default, PyTorch records the information backpropagation needs at every operation, which takes extra memory and time. When only predicting, not training, no_grad turns it off, using fewer resources and running faster.