PyTorch

PyTorch 用三块积木拼出所有深度学习:张量(tensor)存数据、autograd 自动算梯度、nn.Module 组织模型。本文从这三块积木讲起,给出一套可直接改用的训练管线,并覆盖多卡与推理部署;代码基于 torch 2.4 API 编写,在独立虚拟环境安装后即可运行(pip install torch --index-url https://download.pytorch.org/whl/cpu 可装 CPU 版验证)。

一、张量:数据与设备

• 张量 = 多维数组 + dtype + device;形状操作遵循「N 维数组」直觉:

import torch
x = torch.arange(12).reshape(3, 4)      # 3x4
y = x.T                                  # 转置 4x3
print(x + 1, x.sum(dim=0))               # 广播加法;按第 0 维求和

• 广播规则:两个张量从后往前对齐维度,尺寸一致或一方为 1 则可运算(1 会被扩展);

a = torch.zeros(3, 1); b = torch.ones(1, 4)
print((a + b).shape)                     # torch.Size([3, 4]),自动扩展

• dtype 与 device 是玄学重灾区:CUDA 上写死的 CPU 张量、float64 拖慢、.item() 取标量、.detach().cpu().numpy() 转 numpy;混合精度见下文 AMP;
• 设备习惯:model.to(device)、每个 batch 也 .to(device);多卡用 device_ids 或 DistributedDataParallel 自动搬运。

二、自动求导与计算图

• 训练核心咒语:requires_grad=True 的张量参与运算会记录计算图,调用 loss.backward() 把梯度写进每个叶张量的 .grad;
• 三件套防误伤:

w = torch.tensor([2.0, -1.0], requires_grad=True)
loss = (w * w).sum()
loss.backward()
print(w.grad)                 # tensor([4., -2.]) = d(w^2)/dw

• 推理与「不该求导」时:优先 torch.no_grad()(省显存与时间);只阻断某段用 .detach();更新参数后手动清零:optimizer.zero_grad() 或在 backward 前清零,否则梯度跨 batch 累加;
• 直觉:计算图是前向的账本,backward 沿账本反向分摊损失——见 深度学习·反向传播;不想要梯度追踪时(如特征提取、超参扫描)用 requires_grad_(False)。

三、模型搭建:nn.Module

• 一切模型继承 nn.Module,实现 forward;nn.Sequential 拼线性栈,自定义逻辑直接写 forward:

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim=28*28, hidden=128, out=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, out))
    def forward(self, x):
        return self.net(x)

model = MLP()
print(sum(p.numel() for p in model.parameters()))   # 参数总数

• 常用积木:nn.Linear / Conv2d / BatchNorm / Dropout / Embedding / LayerNorm / MultiheadAttention;nn.ModuleList 放模块列表(注意不是 Python list,否则参数不注册);
• 损失函数库 nn.:CrossEntropyLoss(自带 softmax,别在模型里再套一次)、MSELoss、BCEWithLogitsLoss;优化器 torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)。

四、数据管线:Dataset 与 DataLoader

from torch.utils.data import Dataset, DataLoader

class ToyDS(Dataset):
    def __init__(self, n=200):
        import math
        self.xs = torch.rand(n) * 4 - 2
        self.ys = torch.sin(self.xs) + torch.randn(n) * 0.1
    def __len__(self):
        return len(self.xs)
    def __getitem__(self, i):
        return self.xs[i].reshape(1), self.ys[i].reshape(1)

loader = DataLoader(ToyDS(), batch_size=16, shuffle=True, num_workers=0)

• 要点:shuffle 只在训练开;验证/测试不开(shuffle=False);num_workers 在 Linux 多进程加载可提速,Windows/mac 上从 0 开始调;
• 进阶:图/变长数据用自定义 collate_fn(如对序列做 padding 与 attention_mask)。

五、完整训练循环(可直接套用)

import torch, torch.nn as nn
from torch.utils.data import Dataset, DataLoader

model = nn.Sequential(nn.Linear(1, 64), nn.ReLU(), nn.Linear(64, 1))
opt   = torch.optim.AdamW(model.parameters(), lr=1e-2)
loss_fn = nn.MSELoss()
loader = DataLoader(ToyDS(n=300), batch_size=16, shuffle=True)

for epoch in range(200):
    total = 0.0
    for xb, yb in loader:
        opt.zero_grad()
        pred = model(xb)
        loss = loss_fn(pred, yb)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 防爆炸
        opt.step()
        total += loss.item() * len(xb)
    if epoch % 50 == 0:
        print(f"epoch={epoch}  mse={total / len(loader.dataset):.4f}")

• 代码骨架固定为:取 batch → 清梯度 → 前向 → 算损失 → 反向 → 可选梯度裁剪 → 更新;验证时包一层 with torch.no_grad(),用同一 loss 计算验证误差判断过拟合;
• 训练脚本组织:model / data / optimizer / metric 四段 + 每 N 步打印;写一个可复用的 train_one_epoch 函数是所有项目的第一步。

六、训练进阶:AMP、早停与检查点

scaler = torch.amp.GradScaler("cuda")     # 需 GPU;CPU 可不启用
for xb, yb in loader:
    opt.zero_grad()
    with torch.amp.autocast("cuda"):
        loss = loss_fn(model(xb), yb)
    scaler.scale(loss).backward()
    scaler.step(opt)
    scaler.update()

• AMP(自动混合精度):fp16 计算 + fp32 权重,现代 GPU 上吞吐近翻倍;GradScaler 防小梯度下溢;
• 早停:验证 loss 连续 N 个 epoch 不降则停止,并回滚到最优权重;
• 检查点:保存 {"model": model.state_dict(), "opt": opt.state_dict(), "epoch": e} 与最优验证权重两份;载入时先 load_state_dict 再进 eval 模式(model.eval() 影响 Dropout/BN 行为)。

七、多卡与分布式

• 最简多 GPU:model = nn.DataParallel(model)(单机多卡,弃用趋势);生产推荐 DDP:

# 启动:torchrun --nproc_per_node=8 train.py
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
dist.init_process_group("nccl")                  # 每个进程一个 rank
model = DDP(model.to(rank))                      # 每卡一份模型+数据分片

• DDP 语义:每进程一份梯度 → all-reduce 求平均 → 同步更新;batch 大小/学习率按总卡数对应放大(线性缩放规则);
• 超大模型路线:FSDP(分片优化器与参数)、张量并行、流水并行——实操从 Hugging Face accelerate / transformers 起步比自己手搓省力;
• 显存估算:模型权重 + 优化器状态(Adam 约 2× 权重大小)+ 梯度 + 激活;fp16/分片各砍一块(详见 AI 工程化 的量化与推理内存表)。

八、推理与部署

• 推理标准动作:model.eval(),包 torch.no_grad(),model(x);需要确定性可关掉一部分 CUDA 算子;
• torch.compile:一行 model = torch.compile(model),PyTorch 2.x 默认图编译,GPU 上常见 20-50% 提速(首次调用有编译预热);CPU 也可用但收益视算子而定;
• 导出:ONNX(torch.onnx.export)跨框架、TorchScript 保兼容但维护成本高;Hugging Face 生态用 save_pretrained + safetensors 存权重的时代已到;
• 落地部署:服务化框架与 LLM 推理栈见 AI 工程化;大模型微调工具链见 模型微调。

九、踩坑清单

1. 训练模式没切:忘 model.eval() 时 Dropout/BN 依旧生效,验证分数虚低;预测结果对不上再查这处;
2. 梯度没清零:opt.zero_grad() 漏写,梯度跨 batch 累积,loss 曲线怪;
3. CrossEntropyLoss 又套 softmax:双重 softmax 数值不稳且梯度畸变;
4. 设备不一致:模型在 CUDA、输入还在 CPU 直接报错——统一入口 data = data.to(next(model.parameters()).device);
5. .item() 前没 .detach():backward 卡住或显存泄漏(保留计算图);
6. 广播坑:loss.item() 打印前先确认 loss 是标量(形状 [B,1] 时用 .mean());
7. 固定 seed 只 seed 一次:数据顺序、dropout、权重初始化都要在创建前 seed 且 DataLoader 加 generator 才可完全复现;
8. num_workers 拍大数:worker 常驻内存,小数据集反而慢;Windows 上 worker>0 需 if __name__ == "__main__" 保护;
9. 检查点只存模型不存优化器:断点续训时优化器状态(动量/自适应)丢失,曲线异常;
10. AMP 在 CPU/老卡硬开:不支持时静默变 fp32 或报错,先用小网络验证环境。

继续学习:PyTorch 是所有深度模型的引擎,上层应用如 RAG 向量检索与 模型微调 的 LoRA 训练,底层都是本文的循环骨架。