Day 12: 深度学习框架与工程实践

摘要：懂得理论只是第一步，真正的挑战在于把模型跑起来。Day 12 将带你深入 PyTorch 的核心机制（Autograd、Module），揭秘混合精度训练（AMP）如何让显存翻倍，并探讨梯度检查点等工程优化技巧，助你从“调包侠”进阶为“架构师”。

---

1. PyTorch 核心机制

PyTorch 之所以流行，是因为它动态图的设计符合 Python 程序员的直觉。

1.1 Autograd (自动求导)

这是 PyTorch 的魔法引擎。它会记录你对张量做的所有操作，构建一个计算图 (Computational Graph)，然后在反向传播时自动计算梯度。

• Leaf Tensor (叶子节点)：用户创建的张量（如权重 $W$）。requires_grad=True 表示需要对它求导。
• Function (算子)：加减乘除等操作，它们连接了叶子节点。
• Graph (图)：loss.backward() 时，引擎会从 loss 节点出发，沿着图反向走，利用链式法则算出每个叶子节点的梯度 .grad。
• 反向：看着草稿纸的步骤，倒推每个变量对结果的影响（求导）。
• 可视化示例：$y=(a+b)\times c$

  [a] --> (Add) --> [x] --> (Mul) --> [y]
  [b] --^            ^
  [c] ---------------|
  

  打印 y.grad_fn 会看到 <MulBackward>，这就是图上的节点。

1.2 nn.Module 与 State Dictionary

• nn.Module：所有神经网络层的基类。它帮你管理网络中的参数（self.parameters()）。
• state_dict：一个 Python 字典，存着模型所有的参数。
  • 权重：{'layer1.weight': tensor([...])}
  • 优化器状态：除了权重，还需要存优化器的“记忆”（如 Momentum 的动量）。如果不存，恢复训练时动量会丢失（速度归零），导致Loss震荡。
• ckpt / .pth：（Key是层名，Value是张量）。保存/加载模型其实就是保存/加载这个字典。
  • torch.save(model.state_dict(), 'ckpt.pth')

---

2. 混合精度训练 (Mixed Precision Training)

随着模型越来越大，显存和速度成了瓶颈.
随着模型越来越大，显存和速度成了瓶颈。

• FP32 (单精度)：占用 4 字节。精度高，但慢且占地大。
• FP16 (半精度)：占用 2 字节。快且省地，但精度低，容易溢出（太大的数变无穷）或下溢（太小的数变0）。

2.1 AMP (Automatic Mixed Precision)

PyTorch 的 AMP 策略是：该高精度的地方用 FP32，能省的地方用 FP16。

• 它是自动的，但基于预设规则：卷积/矩阵乘法自动转 FP16，Softmax/Loss 自动转 FP32。

2.2 Loss Scaling (损失缩放)

为了解决 FP16 下溢 (Underflow) 问题（梯度太小，FP16 存不下直接变 0）：

• 原理：先放大，后缩小。
  1. Scaling：把 Loss 乘以一个大数（如 $2^{16}$）。这样梯度也跟着放大了，就能被 FP16 记录下来。
  2. Unscaling：更新权重前，再除以 $2^{16}$，变回真实数值。
• 质疑：除回去不又变成 0 了吗？
  • 不会。因为更新权重的步骤通常是在 FP32 下进行的（Optimizer 内部维护 FP32 的权重副本），FP32 能存下极小的数。

💻 代码实践：AMP 详解

import torch
from torch.cuda.amp import autocast, GradScaler

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters())
scaler = GradScaler()  # 1. 创建缩放器，负责管理放大倍数

for input, target in dataloader:
    input, target = input.cuda(), target.cuda()
    
    # 2. 开启自动混合精度上下文
    # 在这个 with 块里的操作，PyTorch 会自动判断是用 FP16 还是 FP32
    with autocast():
        output = model(input)
        loss = criterion(output, target)
    
    optimizer.zero_grad()
    
    # 3. 反向传播：Scale Loss
    # 不直接做 loss.backward()，而是先放大 loss
    # 这样算出来的梯度也是放大的，防止在 FP16 下变成 0
    scaler.scale(loss).backward()
    
    # 4. 权重更新：Unscale -> Update
    # scaler.step 会先尝试把梯度除回去 (Unscale)
    # 如果发现除回去后有 Inf/NaN (溢出)，这一步就会跳过，不更新参数（安全机制）
    # 如果正常，就用 optimizer.step() 更新参数
    scaler.step(optimizer)
    
    # 5. 更新缩放因子
    # 如果这几步都很稳定，scaler 可能会尝试增加放大倍数
    # 如果这步溢出了，scaler 会减小放大倍数
    scaler.update()

---

3. 显存优化黑科技

大模型时代，显存就是金钱。除了 AMP，还有两招必杀技。

3.1 梯度检查点 (Gradient Checkpointing)

• 区分概念：
  • Model Checkpoint：存档（存硬盘）。
  • Gradient Checkpointing：显存优化技术。
• 原理：时间换空间。
  • 正常训练需要保存每一层的中间结果（Activations）用于反向传播，很占显存。
  • 开启此功能后，只保存部分关键节点的中间结果。需要用到其他结果时，临时重新计算一遍前向传播。
  • 效果：显存占用大幅降低（适合跑大模型），但训练速度变慢（多了重计算的时间）。
  • model.gradient_checkpointing_enable()

3.2 梯度累积 (Gradient Accumulation)

• 场景：显存太小，Batch Size 只能开到 32，但你想跑 128 的效果。
• 原理：积少成多。
  • Step 1: 跑 32 个数据，算梯度，不更新参数，把梯度攒起来。
  • Step 2: 再跑 32 个… 攒起来。
  • …
  • Step 4: 攒够了 ($32\times 4=128$)，执行一次 optimizer.step() 更新参数，并清空梯度。
• Mini-batch: 通常指一次 forward 喂进去的数据（这里的 32）。

---

4. 训练流程 SOP 与调试

当你开始跑模型时，不能太“佛系”，要有系统化的意识。

4.1 跑前、跑中、跑后 (SOP)

1. 跑之前：检查数据形状、Label是否对应、学习率设置是否合理。
2. 跑之中 (Monitoring)：
   • Loss：是否在下降？如果是 NaN，通常是梯度爆炸或除零。
   • 显存：nvidia-smi 监控。
   • 速度：如果太慢，用 Profiler 查瓶颈。
3. 跑之后：分析 Bad Case（预测错的样本），看是模型能力问题还是数据标注问题。

4.2 监控工具：TensorBoard vs WandB

• TensorBoard：老牌，本地运行。适合离线调试，隐私性好。
• WandB (Weights & Biases)：行业标准，强烈推荐。
  • 云端同步：手机也能看训练进度。
  • 系统监控：自动记录 GPU 温度、显存、CPU 负载，一眼看出是不是卡在 IO 上。
  • 实验管理：自动记录超参数，方便对比不同实验的效果。

4.3 Profiler (性能分析器)

• 是什么：代码的“体检报告”。
• 作用：告诉你程序每一步花了多少时间。
• 使用时机：不要跑完全程！ Profiler 采样很慢。通常只采样中间的 10 个 Step 就足以发现瓶颈（如 GPU 利用率低、DataLoader 卡顿）。

with torch.profiler.profile(...) as p:
    model(input)
print(p.key_averages().table(sort_by="cuda_time_total"))

---

5. 总结

深度学习不仅仅是设计模型，更是一门系统工程。

• AMP 是现代训练的标配，能白嫖一倍的性能。
• 梯度检查点 让你在有限显存下跑起更大的模型。
• 梯度累积 解决了小显存跑大 Batch 的难题。

掌握这些，你才能真正驾驭大模型训练。

---

参考资料

• PyTorch Automatic Mixed Precision Examples
• Fitting Larger Networks into Memory (OpenAI)