一个可运行的 PyTorch 训练循环

DEEP LEARNING HANDBOOK / 01

让一个小模型学会 y = 2x + 1

一个 CPU 上的教学例子，用来观察梯度更新；它不是真实任务的泛化能力评测。

1. 准备环境和张量

先按 PyTorch 官方安装说明在独立 Python 环境安装与你的系统匹配的版本。此例不需要 GPU、数据下载或额外的 torchvision 包。

输入是四个数，每个样本只有一个特征，所以形状为 [4, 1]。检查 shape、dtype 和 device 是排查张量问题的起点，参见官方张量教程。

2. 完整示例

import torch
from torch import nn

torch.manual_seed(7)
x = torch.tensor([[-1.0], [0.0], [1.0], [2.0]])
y = 2 * x + 1

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

model.train()
for step in range(100):
    optimizer.zero_grad(set_to_none=True)
    prediction = model(x)
    loss = loss_fn(prediction, y)
    loss.backward()
    optimizer.step()
    if step % 20 == 0:
        print(step, float(loss.detach()))

model.eval()
with torch.no_grad():
    print("weight:", model.weight.item())
    print("bias:", model.bias.item())
    print("prediction at x=3:", model(torch.tensor([[3.0]])).item())

3. 一次更新中的五件事

先清理梯度，再做预测、计算损失、反向传播、更新参数。PyTorch 默认累积梯度，因此不清理会改变这个示例的优化行为。具体 API 与训练步骤见官方优化教程。

nn.Linear(1, 1) 是一条直线：一个权重加一个偏置；MSE 衡量预测和目标的平方差。对于本例，合理的检查是权重逐渐接近 2、偏置接近 1，且 x=3 时的预测接近 7。不要把示例的数值写成已经测得的研究结果。

4. 区分推理与评价

model.eval() 切换训练相关层的模式；torch.no_grad() 关闭这个代码块中的梯度记录，二者不是同一个功能。虽然本模型没有 Dropout，仍保留这两步以建立正确习惯，见 Module 官方 API。

训练集损失下降不证明模型对新数据有效。真实任务要先划分训练、验证与测试数据；同一病人、同一来源等关联样本应考虑分组划分，避免信息泄漏。

小练习

把学习率改小，比较相同步数后权重和偏置距离目标有多远；再将标签改为 3*x-2，先预测合理结果，再运行检查。

返回深度学习手册继续阅读过拟合与评价。