aopstudio 的个人博客

学完《动手学深度学习》（D2L） 3.2 的线性回归从零实现，再看 3.3 的线性回归简洁实现时，我一下子有点跟不上：原先自己写出来的前向计算、损失函数和参数更新，怎么突然只剩下四行了？

for X, y in data_iter:
    l = loss(net(X), y)
    trainer.zero_grad()
    l.backward()
    trainer.step()

后面慢慢理解学习才发现，这四行并没有省掉训练中的关键步骤，只是把原来需要手写的部分交给了 PyTorch自动完成，对于开发人员来说更加方便了。下面就顺着 D2L 3.2 的代码，把它们一一找回来。

net：保存参数，负责预测

在 3.2 从零实现线性回归中，线性回归的参数和前向计算是自己写的：

w = torch.normal(0, 0.01, size=(2, 1), requires_grad=True)
b = torch.zeros(1, requires_grad=True)

def linreg(X, w, b):
    return torch.matmul(X, w) + b

w 和 b 是需要学习的参数，linreg 用它们把输入 X 算成预测值。写成 3.3 的形式后，这两件事交给了 nn.Linear：

net = nn.Sequential(nn.Linear(2, 1))

y_hat = net(X)

nn.Linear(2, 1) 内部已经有权重和偏置，相当于自动创建了 w 和 b，并把它们注册为需要求梯度的参数（也就是设置了requires_grad=True）。net(X) 做的仍然是线性回归的前向计算，和手动写出的 torch.matmul(X, w) + b 本质相同，只是不必再把 w、b 作为函数参数传进去。

这里的线性回归其实只有一层 nn.Linear(2, 1)。外面再包一层 nn.Sequential()，是因为实际模型往往由多层组成，之后可以继续往里面添加；这个例子虽然简单，写法已经和更一般的模型保持一致。

所以，net 可以先理解成“模型本身”：它保存参数，也知道如何根据输入给出预测，但它不负责判断预测得好不好，更不知道该怎样修改参数。

loss：根据预测和标签给出这一批数据的误差

3.2 中的平方损失也是手写的：

def squared_loss(y_hat, y):
    return (y_hat - y.reshape(y_hat.shape)) ** 2 / 2

l = squared_loss(linreg(X, w, b), y)

这里的 l 里存着一个批次中每个样本的损失，是个张量，而反向传播计算梯度的对象只能是个标量，所以训练时还要把它们求和：

l.sum().backward()

而 3.3 使用了 PyTorch 已经实现好的均方误差损失：

loss = nn.MSELoss()
l = loss(net(X), y)
...
l.backward()

nn.MSELoss() 默认对损失做平均，因此这里的 l 已经是一个标量，后面可以直接写 l.backward()。

loss 的职责到这里就结束了：它接收预测值和真实标签，产出一个可以衡量当前误差的 l。它不保存模型参数，也不更新参数。

trainer：拿着参数，按指定规则更新它们

在 3.2 中，更新参数的规则是自己写的：

def sgd(params, lr, batch_size):
    with torch.no_grad():
        for param in params:
            param -= lr * param.grad / batch_size
            param.grad.zero_()

这段函数需要把参数列表、学习率和批量大小都传进来。同时注意 3.2 中手写的 SGD 在更新时还要除以 batch_size计算平均值（而 3.3 的 l 已经是平均损失，梯度也就不需要再在 step() 时除以 batch_size）

3.3 则在训练前创建优化器：

trainer = torch.optim.SGD(net.parameters(), lr=0.03)

torch.optim.SGD 表示本次训练使用 SGD 优化器；net.parameters() 把模型中的权重和偏置交给 trainer 管理；lr=0.03 则指定学习率。之后调用：

l.backward()
trainer.step()

l.backward() 会计算当前损失对模型参数的梯度，并把结果写入参数的 .grad 属性；trainer.step() 再读取这些 .grad，按 SGD 的规则原地修改参数。若把优化算法换成 Adam，模型 net 和损失 loss 都不必改，只需把 torch.optim.SGD() 换成 torch.optim.Adam()。

这样看，3.2 里的 sgd([w, b], lr, batch_size) 和 3.3 里的 trainer.step() 对应的是同一件事。区别只是参数和更新规则不再每一批都手动传入，而是在创建优化器时保存好了。

一次训练到底发生了什么

现在回到最开始的训练循环：

l = loss(net(X), y)
trainer.zero_grad()
l.backward()
trainer.step()

第一行依次调用了 net 和 loss：先得到预测值，再得到当前批量的损失 l。

l.backward() 从当前损失出发，计算它对相关模型参数的梯度，并把结果放到这些参数的 .grad 属性中。至于 PyTorch 为什么能沿着计算过程求出梯度，是前一节自动求导讨论的内容，此处不再赘述。

trainer.step() 才会使用 .grad 和创建时给定的 SGD、学习率来更新参数。它不负责计算损失，也不负责决定梯度；它只执行优化规则。

为什么 PyTorch 不将 l.backward() 和 trainer.step() 合成一个操作？

backward() 与 step() 分开不是多写了一行。前者回答“当前损失对参数的梯度是多少”，后者回答“按什么规则利用这些梯度更新参数”，这是两项不同的职责。

而且，在高级计算中，两者之间还可以插入别的操作，使得计算更加灵活，例如梯度裁剪：

l.backward()
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
trainer.step()

所以它们保持分开，训练时才有足够的调整空间。

为什么是 trainer.zero_grad()，不是 net.zero_grad()

3.2 的 sgd 函数末尾写的是：

param.grad.zero_()

这里的 zero_() 是对某一个参数的梯度张量做原地操作。3.3 中不需要再把这个动作写进手写的 sgd，因为优化器提供了：

trainer.zero_grad()

这里先补一个容易被忽略的细节：zero_grad() 的目的，是清掉上一个批量留下的梯度，避免下一次 backward() 把新旧梯度累加起来。按 PyTorch 目前的默认设置，它通常把 .grad 设为 None，不一定是逐元素填成 0；对这一节来说，关键是旧梯度不再参与下一次更新。

在本节这个简单模型里，net.zero_grad() 和 trainer.zero_grad() 处理的是同一批参数：net 的全部参数都交给了这个 trainer。理论上来说，调用 net.zero_grad() 其实也能清掉模型参数的旧梯度，和trainer.zero_grad()产生的结果是一样的，但两者的出发点不同：

• net.zero_grad() 面向模型拥有的参数；
• trainer.zero_grad() 面向这个优化器管理的参数。

trainer.zero_grad() 之后紧接着要调用的是 trainer.step()，两句都针对优化器实际要更新的那批参数，因此教程写成 trainer.zero_grad() 更连贯。以后如果一个模型的不同部分交给不同优化器管理，两者涉及的参数范围就可能不同；这时更要看清自己希望操作的是模型全部参数，还是某个优化器负责的参数。

为什么是 trainer.step()，而不是 net.step()？

net 知道模型的结构、参数和前向计算，但它并不知道训练时选的是 SGD、Adam 还是 AdamW，也不知道学习率等优化超参数。它无法替优化器决定更新规则。

trainer 在创建时保存了这些规则，也持有参数对象的引用，所以 trainer.step() 可以根据上述更新规则来直接修改 net 中的参数。

为什么不让所有操作都由 trainer 托管？

因为三者掌握的信息不同。net 负责把输入变成预测；loss 负责把预测和标签变成损失；得到具体的 l 后，才能调用 l.backward()。而 trainer 的职责是管理参数的更新规则。

让优化器直接包办前向计算和损失计算，会把模型结构、损失函数和优化算法绑在一起，不利于排查问题和维护。现在的分工使得同一个模型能配不同损失或优化器，也让训练循环仍然看得见每一步发生了什么。

小结

D2L 3.2 把线性回归的每个零件露在外面：w、b、linreg、squared_loss 和 sgd。3.3 没有改变训练的基本顺序，只是把模型收进 net，把误差计算收进 loss，把参数更新收进 trainer。

以后再看到类似的训练循环，可以按同样的顺序去读：net 先预测，loss 给出误差，backward() 把梯度写到参数上，trainer 再决定如何用这些梯度更新参数。模型可以换成卷积网络，损失也可以换成交叉熵，但这条分工主线不会变。把这四个位置看明白，简洁实现也就不再像是把训练过程藏起来了。