PyTorch 入门

PyTorch 入门

本节摘要:你已经从活塞和曲轴造好了引擎,现在来学大家真正在开的那台。你有一个能跑的迷你框架:Linear 层、ReLU、Dropout、BatchNorm、Adam、DataLoader、训练循环,它能用纯 Python 在圆形分类上训一个 4 层网络。但它在同一问题上比 PyTorch 慢约 500 倍。你的框架一次处理一个样本、嵌套 Python 循环;PyTorch 把同样运算派发给跑在 GPU 上的优化 C++/CUDA 内核。单块 NVIDIA A100 上,PyTorch 约 6 小时能在 ImageNet(128 万张图)上训完 ResNet-50(2560 万参数);你的框架同样任务大约要 3000 小时——如果没先撑爆内存的话。速度还不是唯一差距:你的框架没 GPU 支持、没自动微分(每个模块的 backward() 都是你手写的)、没序列化、没分布式训练、没混合精度、除了 print 没法调试梯度流。PyTorch 把这些缺口全填上,而且用的是你已经建好的同一套心智模型:Module、forward()、parameters()、backward()、optimizer.step()——概念一一迁移,语法几乎相同,差别只在 PyTorch 把十年系统工程包进了你从零设计的同一套接口背后。本节用 PyTorch 在 MNIST 上从零训一个 3 层 MLP,讲清张量、autograd、nn.Module 与标准训练循环,并与你的迷你框架逐行对照。

学习目标

阅读完本节,你应当能够:

1. 用 PyTorch 的 nn.Module、nn.Sequential、autograd 构建并训练神经网络。
2. 使用 PyTorch 张量、GPU 加速与标准训练循环(zero_grad、forward、loss、backward、step)。
3. 把你从零实现的迷你框架组件转换为 PyTorch 等价物。
4. 在同一任务上剖析并比较纯 Python 框架与 PyTorch 的训练速度。

一、问题与直觉

你有一个能跑的迷你框架——它在纯 Python 上训练一个 4 层圆形分类网络,但同一问题比 PyTorch 慢约 500 倍。

你的框架一次处理一个样本、嵌套 Python 循环;PyTorch 把同样运算派发给跑在 GPU 上的优化 C++/CUDA 内核。单块 A100 上 PyTorch 约 6 小时训完 ImageNet 上的 ResNet-50,你的框架大约要 3000 小时——如果没先撑爆内存的话。

速度还不是唯一差距:你的框架没 GPU 支持、没自动微分(每个模块的 backward() 都手写)、没序列化、没分布式训练、没混合精度、除了 print 没法调试梯度流。PyTorch 把这些缺口全填上,而且用的是你已经建好的同一套心智模型:Module、forward()、parameters()、backward()、optimizer.step()。概念一一迁移,语法几乎相同,差别只在 PyTorch 把十年系统工程包进了你从零设计的同一套接口背后。

PyTorch 为何赢了

2015 年 TensorFlow 要求你先定义静态计算图再运行:搭图、编译、喂数据,调试意味着盯着图可视化,改架构意味着从零重建图。PyTorch 2017 年以不同理念登场:即时执行(eager execution)——你写 Python,它立刻运行。y = model(x) 此刻就真的算出 y,而不是「往图里加一个稍后算 y 的节点」。这意味着标准 Python 调试工具都能用:print() 能用、pdb 能用、forward 里的 if/else 都能用。

到 2020 年市场已经表态:PyTorch 在 ML 研究论文里的份额从 7%(2017)涨到 75% 以上(2022)。Meta、Google DeepMind、OpenAI、Anthropic、Hugging Face 都把 PyTorch 作为主力框架。TensorFlow 2.x 转向即时执行作为回应——变相承认 PyTorch 的设计是对的。教训:开发者体验会复利,一个慢 10% 但调试快 50% 的框架每次都赢。

张量(Tensor)

张量是带三个关键属性的多维数组:shape、dtype、device。

import torch

x = torch.zeros(3, 4)           # shape: (3, 4), dtype: float32, device: cpu
x = torch.randn(2, 3, 224, 224) # 2 张 RGB 图的 batch, 224x224
x = torch.tensor([1, 2, 3])     # 从 Python list 创建
​

shape 是维度:标量是 (),向量是 (n,),矩阵是 (m, n),图像 batch 是 (batch, channels, height, width)。dtype 控制精度与内存:

device 决定计算在哪发生。所有运算都要求张量在同一设备——这是新手最常撞的 PyTorch 错误:RuntimeError: Expected all tensors to be on the same device,修正办法是计算前把所有张量挪到同一设备。重塑是常量的——只改元数据,不动数据:x.view()、x.reshape()、x.permute()、x.unsqueeze()、x.squeeze()。

Autograd

你的迷你框架要求每个模块手写 backward(),PyTorch 不用。它把张量上的每次运算记录进有向无环图(计算图),然后反向遍历该图自动算梯度。

关键差别:PyTorch 用基于「磁带」的自动微分,前向时每次运算追加到「磁带」,调用 .backward() 时反向回放磁带。

x = torch.randn(3, requires_grad=True)
y = x ** 2 + 3 * x
z = y.sum()
z.backward()
print(x.grad)  # dz/dx = 2x + 3
​

autograd 三规则:

1. 只有 requires_grad=True 的叶张量才累积梯度;
2. 梯度默认累积——每次反向前要调 optimizer.zero_grad();
3. torch.no_grad() 关闭梯度跟踪(评估时用)。

nn.Module

nn.Module 是 PyTorch 里每个神经网络组件的基类,你在第 10 节已经建过这个抽象。PyTorch 版多了自动参数注册、递归模块发现、设备管理、state_dict 序列化。

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = self.layer1(x)
        x = self.relu(x)
        x = self.layer2(x)
        return x
​

在 __init__ 里把 nn.Module 或 nn.Parameter 赋给属性时,PyTorch 自动注册它。model.parameters() 递归收集每个已注册参数——这就是你不必像迷你框架那样手工收集权重的原因。

关键积木:

损失函数与优化器

PyTorch 内置你造过的一切的生产级版本。注意:CrossEntropyLoss 内部合并了 LogSoftmax + NLLLoss,要传原始 logits 而非 softmax 输出,这是常见的静默错误来源。优化器方面,CNN 仍常用 SGD + 动量(lr=0.01~0.1),Adam 是默认起点(lr=1e-3),AdamW 用于 Transformer 与微调(lr=1e-4~1e-3)。

训练循环

每个 PyTorch 训练循环都遵循同一套五步模式,你在第 10 节已经会了。

典范模式:

for epoch in range(num_epochs):
    model.train()
    for inputs, targets in train_loader:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
​

batch 循环里就这五行。这五行训出了 GPT-4、Stable Diffusion、LLaMA。架构会变、数据会变,这五行不变。

Dataset 与 DataLoader

PyTorch 的 Dataset 是带 __len__ 与 __getitem__ 两方法的抽象类,DataLoader 给它加上批处理、打乱、多进程加载。num_workers=4 起 4 个进程在 GPU 训当前 batch 时并行加载下一批,在磁盘受限任务(大图、音频)上单这一项就能让训练速度翻倍。

GPU 训练

把模型挪到 GPU 就是 model = model.to(device)——它递归把每个参数和缓冲区搬到 GPU。然后训练时每批挪一次:inputs, targets = inputs.to(device), targets.to(device)。混合精度在现代 GPU(A100、H100、RTX 4090)上把显存减半、吞吐翻倍——前向/反向用 float16,主权重保持 float32,用 torch.amp.autocast 与 GradScaler 包住。

对比:迷你框架 vs PyTorch vs JAX

二、从零实现

用纯 PyTorch 原语在 MNIST 上训一个 3 层 MLP——不用高级封装、不用 torchvision.datasets,我们自己下载并解析原始数据。完整代码见原课程 phases/03-deep-learning-core/11-intro-to-pytorch/code/ 相应文件。

Step 1:从原始文件加载 MNIST

MNIST 是 4 个 gzip 文件:训练图(60000×28×28)、训练标签、测试图、测试标签。我们下载并用 struct 解析二进制格式,转成 float32 并归一化到 [0, 1]。

Step 2:定义模型

3 层 MLP:784 → 256 → 128 → 10,ReLU 激活,Dropout 正则。输出层产出 10 个原始 logits(每类一个),不加 softmax——CrossEntropyLoss 内部处理。参数量:784×256 + 256 + 256×128 + 128 + 128×10 + 10 = 235,146,按现代标准极小,GPT-2 small 都有 1.24 亿。

class MNISTModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 256), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(256, 128), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(128, 10),
        )

    def forward(self, x):
        return self.net(x)
​

Step 3:训练循环

典范的 forward-loss-backward-step 模式。评估时用 model.eval() 与 torch.no_grad()——后者关闭 autograd,省显存、加速推理,不加它 PyTorch 会建一个你根本不用的计算图。

Step 4:把一切接起来

下载 MNIST、建 Dataset/DataLoader、实例化模型、用 Adam(lr=1e-3)与 CrossEntropyLoss 训 10 轮,最后 torch.save(model.state_dict(), "mnist_mlp.pt") 存权重。预期 10 轮后测试精度约 97.8%;CPU 训练约 30 秒,GPU 约 5 秒,你的迷你框架同样架构约 45 分钟。

三、框架对比

迷你框架 vs PyTorch 速查

接口几乎一致,差别全在引擎盖下。

存取模型

torch.save(model.state_dict(), "model.pt")

model = MNISTModel()
model.load_state_dict(torch.load("model.pt", weights_only=True))
model.eval()
​

永远存 state_dict()(参数字典)而非模型对象——存对象用 pickle,重构代码就坏;state_dict 可移植。

学习率调度

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
for epoch in range(10):
    train_one_epoch(model, train_loader, criterion, optimizer, device)
    scheduler.step()
​

PyTorch 内置 15+ 调度器:StepLR、ExponentialLR、CosineAnnealingLR、OneCycleLR、ReduceLROnPlateau,都插同一个优化器接口。

四、可复用产物

本节产出两个工件(位于原课程 outputs/):

• prompt-pytorch-debugger.md:一个提示,诊断常见 PyTorch 训练失败。
• skill-pytorch-patterns.md:一份 PyTorch 训练模式的技能参考。

五、练习

1. Easy:加批归一化——在每个 Linear 后(激活前)插 nn.BatchNorm1d,比较测试精度与训练速度,它应在更少轮次内达到 98%+。
2. Medium:实现学习率查找器——一轮训练里 LR 从 1e-7 指数升到 1.0,画损失 vs LR,最优 LR 在损失开始爬升之前,用它给 MNIST 模型挑更好的 LR。
3. Hard:移植到 GPU 配混合精度——加 torch.amp.autocast 与 GradScaler,测有/无混合精度的吞吐(样本/秒),A100 上预期约 2 倍提速。

本节要点回顾

1. PyTorch 把十年系统工程包进你从零设计的同一套接口:Module、forward()、parameters()、backward()、optimizer.step() 一一迁移,语法几乎相同,但快约 500 倍。
2. 即时执行是它赢的原因:写 Python 立刻运行,标准调试工具都能用,开发者体验复利;2017 到 2022 年研究份额从 7% 涨到 75%+。
3. 张量三属性:shape、dtype、device;所有运算要求张量在同一设备,这是新手最常撞的错。
4. autograd 基于磁带:前向追加、反向回放,梯度默认累积故每次反向前要 zero_grad,只有 requires_grad=True 的叶张量累积梯度。
5. nn.Module 自动注册参数:赋值即注册,model.parameters() 递归收集,不必像迷你框架手工 gather。
6. 训练循环五步永恒:zero_grad → forward → loss → backward → step,这五行训出了 GPT-4、Stable Diffusion、LLaMA。
7. 存 state_dict 而非模型对象:可移植,pickle 重构就坏。
8. 评估必须 model.eval() + torch.no_grad():前者切 Dropout/BatchNorm 行为,后者关 autograd 省显存。
9. 混合精度 halve 显存、翻倍吞吐:float16 前向/反向,float32 主权重,A100/H100 标配。

下一节,我们看 JAX——它编译纯函数而非即时修改变量,这是 Google DeepMind 训 Gemini、Anthropic 训 Claude 的选择,也彻底改变你对深度学习的思维方式。