PyTorch：从搭积木到构建系统

发表于 更新于 分类于 Waline： 本文字数： 19k 阅读时长 ≈ 1:07

1. 从零搭建一个 PyTorch 模型

在正式讨论 PyTorch 之前，我们先从一个最小可用训练流程入手：用最少的代码搭建并训练一个小型神经网络。通过这个过程来快速建立对 PyTorch 的整体认知：数据从哪里来？如何流入模型？损失如何计算？梯度如何反向传播？参数又是如何被更新的？ 后面章节会对这些环节再做细致拆解，我们先做一个“总览式体验”。

1.1 整体思路：从数据到参数更新

我们开始搭建的最小神经网络包含四个核心阶段：

  1. 数据准备：把原始数据转换为张量，并按批次（batch）组织；
  2. 模型构建：用 nn.Module 定义前向计算逻辑；
  3. 损失与优化器：定义“好坏标准”和“如何更新参数”的规则；
  4. 训练循环：反复执行前向、反向和参数更新。

一个典型的 PyTorch 训练循环大致如下：

1
2
3
4
5
y = model(x)                 # 前向传播：x → 预测 y
loss = criterion(y, target)  # 损失计算：预测y 与 真实y 的误差
loss.backward()              # 反向传播：沿计算图求梯度
optimizer.step()             # 参数更新：根据梯度调整权重
optimizer.zero_grad()        # 清除旧梯度，防止累加
flowchart LR
    A[Dataset
样本定义] --> B[DataLoader
批量加载/打乱]
    B --> C[Model Forward
前向传播]
    C --> D[Loss Function
损失计算]
    D --> E[Backward
反向传播
梯度求导]
    E --> F[Optimizer Step
参数更新]

    %% 反向路径标注（虚线）
    E -.-> C
图1-1 PyTorch 训练数据流动路径图

1.2 构建 Tensor 与参数：从张量开始

在 PyTorch 中，一切计算都围绕 Tensor（张量） 展开。输入数据、模型参数、中间激活值，本质上都是张量。

1
2
3
4
5
6
7
8
9
10
11
12
import torch

# 输入张量：4 个样本，每个样本 10 个维度特征
x = torch.randn(4, 10)

# 权重矩阵：把 10 维输入投影到 5 维输出
W1 = torch.randn(10, 5, requires_grad=True)
# 偏置向量：对应 5 个输出神经元
b1 = torch.zeros(5, requires_grad=True)

# 前向计算：线性变换 + ReLU
h = torch.relu(x @ W1 + b1)   # 形状 [4, 5]

这里有个关键点，requires_grad=True 意味着 需要PyTorch追踪这个张量的梯度，它会被视为可训练参数，一旦参与运算，这些张量背后会动态构建一张计算图，为之后的 backward() 做准备。这一小节先感受“张量 + 运算 = 可求导的计算”，详细的张量机制会在第 2 章展开。

1.3 定义一个两层 MLP：模型即模块

在 PyTorch 中，模型通常通过继承 nn.Module 定义。这样做的好处是参数会被自动注册，便于优化器管理，而且前向逻辑集中在 forward() 中，结构清晰，还可以通过模块组合构建复杂网络。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import torch.nn as nn
import torch.nn.functional as F

class SimpleMLP(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        # 第 1 层：线性变换（输入 → 隐藏层）
        self.fc1 = nn.Linear(in_dim, hidden_dim)
        # 第 2 层：线性变换（隐藏层 → 输出层）
        self.fc2 = nn.Linear(hidden_dim, out_dim)

    def forward(self, x):
        """
        前向传播：说明输入 x 如何一步步变成输出 y
        每次调用 model(x) 时，PyTorch 会自动执行这里的逻辑
        """
        x = F.relu(self.fc1(x))  # 全连接 + ReLU 激活
        x = self.fc2(x)          # 输出层线性变换
        return x

# 实例化模型：输入 784 维（28x28 像素）→ 隐层 256 → 输出 10 类
model = SimpleMLP(784, 256, 10)
print(model)

第 3 章会专门讨论 nn.Module 的模块化设计理念，这里先把它当作“定义模型结构的标准方式”。

1.4 Dataset 与 DataLoader：数据如何流入模型

PyTorch 把“数据的逻辑组织”和“数据的物理加载”明确拆开：

  • Dataset 负责：给定一个索引 idx，返回“第 idx 个样本是什么”；
  • DataLoader 负责：如何按批次、多进程地取样本。

1.4.1 Dataset：数据的逻辑视图

以下示例使用 torchvision.datasets.MNIST，它已经实现了标准的 Dataset 接口：

1
2
3
4
5
6
7
8
9
10
11
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 数据预处理：转为 Tensor 并归一化到 [-1, 1]
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

# 训练集
train_data = datasets.MNIST(root="./data", train=True, download=True, transform=transform)

# 测试集
test_data = datasets.MNIST(root="./data", train=False, download=True, transform=transform)

1.4.2 DataLoader：批量加载

1
2
3
4
# shuffle=True:每个 epoch 打乱数据顺序，提升泛化
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)

test_loader = DataLoader(test_data, batch_size=128, shuffle=False)

这样，我们在训练循环中就可以直接迭代 train_loader，而不用关心底层是如何一条条读文件的。

1.5 完整训练循环与 MNIST 案例

前面我们新建了 SimpleMLP 模型和加载了一个训练集 train_loader 与测试集 test_loader。接下来把它们拼成一个完整的训练流程。

1.5.1 训练循环的三步核心

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import torch.optim as optim

criterion = nn.CrossEntropyLoss()              # 多分类交叉熵
optimizer = optim.SGD(model.parameters(), lr=0.01)

for epoch in range(3):
    for batch_x, batch_y in train_loader:
        # ① 前向传播：构建计算图并得到预测
        batch_x = batch_x.view(batch_x.size(0), -1)  # [B, 1, 28, 28] → [B, 784]
        pred = model(batch_x)
        loss = criterion(pred, batch_y)

        # ② 反向传播：沿计算图求梯度
        optimizer.zero_grad()   # 清空旧梯度
        loss.backward()         # 自动计算 ∂loss/∂参数

        # ③ 参数更新：根据梯度调整权重
        optimizer.step()

    print(f"Epoch [{epoch+1}/3], Loss: {loss.item():.4f}")

1.5.2 在 MNIST 上跑通一个分类器

下方代码是把前面的所有组件整合在一起，训练一个 MNIST MLP 分类器，并在测试集上做简单评估：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
# 省略 import，同前面

# 1. 数据准备（Dataset + DataLoader）
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_data = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
test_data = datasets.MNIST(root="./data", train=False, download=True, transform=transform)

train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=128, shuffle=False)

# 2. 模型定义
class SimpleMLP(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, out_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleMLP(784, 128, 10)

# 3. 损失 & 优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 4. 训练
epochs = 5
for epoch in range(epochs):
    model.train()
    running_loss = 0.0

    for batch_x, batch_y in train_loader:
        batch_x = batch_x.view(batch_x.size(0), -1)

        pred = model(batch_x)
        loss = criterion(pred, batch_y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch [{epoch+1}/{epochs}] - Loss: {running_loss / len(train_loader):.4f}")

# 5. 测试
model.eval()
correct, total = 0, 0
with torch.no_grad():
    for images, labels in test_loader:
        images = images.view(images.size(0), -1)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"\n测试集准确率: {100 * correct / total:.2f}%")

在普通 CPU 或单卡 GPU 上，训练 5 个 epoch 通常可以达到约 96%–97% 的测试准确率。这表示一个“看似朴素”的最小工作流，已经完整覆盖了 PyTorch 的核心计算过程：张量 → 计算图 → 自动微分 → 参数更新。

2. 张量：PyTorch 的核心机制

在 PyTorch 中，有一个几乎贯穿所有机制的核心概念——Tensor（张量）。无论是输入数据、模型参数、激活值，还是计算图中的中间结果，在 PyTorch 的世界中都以张量的形式存在。更重要的是，在 PyTorch 里，张量不只是“存数据的数组”，它还携带了：

  • 是否需要梯度（requires_grad）；
  • 属于哪个设备（CPU / GPU / MPS / …）；
  • 与其它张量的计算关系（计算图上的节点）。

这使得 “张量 + 运算” 不仅完成数值计算，同时也为自动微分系统提供了基础信息。下面我们将在“最小训练循环”的基础上，系统理解 PyTorch 张量和 NumPy 的区别、广播与视图的机制、原地操作的风险，以及设备切换的规范用法。

2.1 Tensor vs NumPy：相似外表下，不同的“内核”

很多人第一次见到 torch.Tensor 的 API 时会有种感觉：这不就是 NumPy 吗？ 确实，在“接口层面”两者很像：都支持加减乘除、索引、切片、广播等。但它们的设计目标却是截然不同：

维度 NumPy 数组 PyTorch Tensor
核心用途 通用数值计算 可微分计算 + 深度学习
执行模式 执行即结束 执行即构图（Define-by-Run）
自动求导 不支持 内置 Autograd
设备支持 CPU CPU / GPU / MPS / XLA / …

2.1.1 NumPy：计算完就结束

1
2
3
4
5
6
7
import numpy as np

x = np.ones((2, 2))
y = (x + 2) ** 2
print(y)
# [[9. 9.]
#  [9. 9.]]

运算执行完，结果就“落地”为一个新的数组；NumPy 不会记录中间计算步骤，也不会保留计算图，更不会去算梯度。

2.1.2 PyTorch：计算的同时在“画图”

对比一下 PyTorch：

1
2
3
4
5
6
7
import torch

x = torch.ones((2, 2), requires_grad=True)
y = (x + 2) ** 2
print(y)
# tensor([[9., 9.],
#         [9., 9.]], grad_fn=<PowBackward0>)

注意输出里的 grad_fn=<PowBackward0> 意味着PyTorch 不只是算出了结果 9，还记录了“这个结果是由一个幂运算得到的”，并且这个幂运算的输入又来自“加法”等操作，这些信息会组成一张计算图，在之后调用 y.backward() 时被 Autograd 引擎用来自动求梯度。

换句话说：在 PyTorch 中，Tensor = 数据 + 梯度信息 + 设备信息 + 计算图中的拓扑位置。

2.2 广播与视图：高效计算的机制

PyTorch 的高效不仅来自 GPU，更来自对 广播（broadcasting） 和 视图（view） 的精细设计。这两者是理解“为什么某些操作几乎不占内存但能完成复杂计算”的关键。

2.2.1 广播：以最少存储完成最大计算表达

广播允许形状兼容的张量在运算过程中“逻辑扩展”，而不真的复制数据。

1
2
3
4
x = torch.randn(3, 5)
b = torch.randn(5)     # 形状 [5]

y = x + b              # 自动广播 b → [3, 5]

这里 b 并不会复制三份，PyTorch 会在内部通过“广播规则”让它在运算时表现得像 [3, 5]，但底层依然只存了一份 [5]。这可以显著节省内存、提升效率。我们可以用一个示意图来理解广播过程：

flowchart LR
    classDef tensor fill:#E8F1FD,stroke:#4A90E2,stroke-width:2px,rx:8,ry:8,color:#111
    classDef op fill:#fff,stroke:#AAB2BD,stroke-width:1.5px,rx:8,ry:8,color:#111,font-size:12px

    A[张量 x
形状：3×5]:::tensor
    B[向量 b
形状：5]:::tensor
    C[广播视图：
b 逻辑扩展为 3×5]:::op
    D[结果 y = x + b
形状：3×5]:::tensor

    A --> C
    B --> C
    C --> D
图 2-1 广播机制示意图
#### 2.2.2 视图（view）：共享内存的张量变形 view() 等操作可以在 **不复制数据** 的前提下，改变张量的形状，这类操作返回的是**同一块底层存储的不同“视图”**。 
1
2
x = torch.randn(4, 4)
y = x.view(2, 8)   # y 和 x 共用同一段内存
特点是修改 y 的数据，会反映到 x 上（只要没打断计算图等）；view() 要求底层内存是连续的，否则会报错，这时可以用 reshape() 兜底（必要时会复制数据）。

小结：广播“逻辑扩展维度”，view“逻辑改变形状”，两者共同实现 “表达复杂运算，但尽量不复制数据” 的目标。

2.3 原地操作（in-place）：是节省内存，但可能带来问题

PyTorch 中所有带下划线的方法，如 x.add_()、x.mul_()、x.relu_()、x.copy_()都是原地操作（in-place），会直接在原内存上修改数据，不再创建新的张量。从数值计算的角度看，这似乎是“更省内存”的做法，但在 自动微分体系（Autograd） 中，它们经常是出问题的点。

2.3.1 为什么 in-place 操作风险很大？

自动微分的本质是：反向传播需要依赖前向传播时的中间结果。如果某个中间结果在 forward 之后被原地改掉，而这个中间结果又是反向传播中需要用到的，那 Autograd 在 backward 时就会发现“我需要的值没了/被覆盖了”，从而报错。看一个例子：

1
2
3
4
5
x = torch.ones(3, requires_grad=True)
y = x + 2
y.relu_()  # 对 y 原地 ReLU
z = y.mean()
z.backward()

在某些场景下，我们可能会看到类似错误：

1
2
RuntimeError: one of the variables needed for gradient computation
has been modified by an inplace operation

这里的关键点并不是“所有 in-place 操作必定会报错”，而是如果这个张量在反向传播中是“必需的中间值”，而我们对它做了原地修改，就可能破坏计算图的一致性。

PyTorch 内部会给参与计算图的 Tensor 维护一个“版本号”，原地修改时版本号会变化。如果 Autograd 发现“反向传播时需要的版本”和“现在的版本”对不上，就会抛出上述错误。

2.3.2 实践经验：训练时尽量避免 in-place

在训练阶段，尤其是对 requires_grad=True 的张量，尽量不要使用 in-place 操作（带下划线的方法）。除非我们非常清楚 Autograd 的细节，并确定这个张量不会再被反向用到，否则宁愿多占一点内存，用非 in-place 版本（如 y = y.relu()）更安全。

2.4 计算设备：CPU / GPU / MPS 之间的迁移

PyTorch 的一大优势是：同一套 Tensor API 可以在不同设备上无缝运行。你可以很自然地把一个模型从 CPU 挪到 GPU，而无需大改代码。

2.4.1 把张量移动到指定设备

常见写法：

1
2
3
4
5
6
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

x = torch.randn(3, 3)
x = x.to(device)              # 显式迁移
# 或者在创建时就指定设备
y = torch.randn(3, 3, device=device)

模型同理：

1
model = SimpleMLP(784, 256, 10).to(device)

训练时只要确保输入张量 .to(device)和模型 .to(device)，就能在 GPU 上完成前向和反向。

2.4.2 跨设备张量无法直接运算

下面的代码会报错：

1
2
3
4
a = torch.randn(2, 2).to("cpu")
b = torch.randn(2, 2).to("cuda")

a + b   # ❌ 报错：不同设备的 Tensor 无法直接运算

规则可以简单记为任何参与同一运算的 Tensor 必须在同一个设备上。所以我们需要：

1
2
b_cpu = b.to("cpu")
c = a + b_cpu   # ✅ 都在 CPU 上

2.4.3 一个简单的 GPU 加速示例

下面的例子我们可以感受大矩阵乘法在 GPU 上的速度差异（在支持 CUDA 的环境中）：

1
2
3
4
5
device = "cuda" if torch.cuda.is_available() else "cpu"

x = torch.randn(10000, 10000, device=device)
y = x @ x      # 大矩阵乘法
print(y.device)

在 CPU 上，这个运算可能要好几秒甚至更久，而在 GPU 上，通常会快一个数量级。实际工程中，大部分深度学习训练都会把模型和大张量迁移到 GPU/MPS 上。

小结：张量是“PyTorch”的最小单元
本章从多个角度回答了“为什么 Tensor 是 PyTorch 的核心机制”这一关键问题。首先，与 NumPy 相比，Tensor 并不仅仅是一个多维数组，而是专为自动微分和多设备计算而设计的计算单元。它内置了 requires_grad 与计算图机制，使得任何运算都能被自动记录并参与反向传播；同时通过 device 属性，可以在 CPU、GPU、MPS 等不同后端之间无缝切换。

Tensor 的广播和视图机制使复杂运算无需复制数据即可完成。广播允许在形状兼容的情况下以逻辑方式扩展张量；视图则在共享底层存储的前提下改变其形状，使得许多高维操作既高效又节省内存。

在自动微分体系中，原地操作则需要格外谨慎。它可能覆写反向传播所依赖的中间结果，从而导致梯度计算错误。因此在训练阶段，应尽量避免使用 in-place 操作，除非明确确认它不会破坏计算图。

最后，设备迁移是性能优化中的基础能力。使用 to(device) 可以以统一方式将张量移动到指定硬件，而不同设备上的张量无法直接进行算术运算，这种约束确保了计算的确定性与安全性。

3. 模块化设计：构建可复用的神经网络组件

在前两章中，我们已经看到 nn.Module 是 PyTorch 中“定义一个模型”的基本方式。但在真实项目中，一个模型往往并不是一个单一的大类，而是由许多小模块拼接、嵌套、组合而成的。

本章的目标是系统回答这些问题：

  • 为什么深度学习框架需要“模块化”？
  • 为什么 PyTorch 的 nn.Module 不是简单的“把几层写在一起”？
  • 怎样让模型像乐高积木一样可组合、易扩展？
  • 如何用 Sequential、自定义模块、子模块树构建复杂网络？
  • 为什么 forward() 中只能写前向逻辑，而不能写参数初始化？

3.1 为什么深度学习模型需要模块化？

一个现代神经网络本质上就是由许多运算层堆叠起来的系统：线性层（Linear）、卷积层（Conv2d）、激活函数（ReLU）、正则化层（Dropout）、残差模块（Residual Block）、Transformer Block（Self-Attention + FFN + Norm）、……

这些“组件”本质上都是有输入、有输出、有内部参数（或无参数），可以组合嵌套的模块。模块化带来的 3 个核心优势：

  1. 参数自动注册，优化器可见
    任何在 nn.Module 中定义的 nn.Parameter 会自动加入参数表中，优化器能一并管理。
  2. 可嵌套的层级结构（module tree）
    一个模型可以由子模块组成，子模块又可以继续包含模块，递归嵌套。
  3. 便于保存 / 加载 / 调试 / 部署
    所有模块都会自动记录自身的结构、参数、缓冲区等，整个网络结构天然具备可序列化性。

这些能力让深度学习模型像搭积木一样灵活，而不需要写成“一个超级长的 forward 函数”。

3.2 神经网络像是一棵模块树（module tree）

理解 PyTorch 模型最核心的概念是：一个模型 = 一棵由 nn.Module 组成的树（Module Tree）。例如下面的模型结构：

1
2
3
4
5
6
7
8
9
10
class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = SimpleMLP(784, 256, 128)
        self.head = nn.Linear(128, 10)

    def forward(self, x):
        x = self.backbone(x)
        x = self.head(x)
        return x

这里Classifier 是根节点，backbone 和 head 是子节点，SimpleMLP 内部又有 fc1、fc2 这些叶子模块。

3.3 nn.Sequential：最轻量的积木拼接方式

当模型结构是按顺序执行的（链式结构），nn.Sequential 是最简洁的写法。例如一个 3 层前馈网络：

1
2
3
4
5
6
7
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

调用 model(x) 会依次执行每个层的 forward。那么什么情况下会使用Sequential？

  • 模型是纯顺序结构；
  • 不需要分支、残差、跳连等复杂逻辑；
  • 没有共享参数的需求；
  • 不需要在 forward 中写复杂逻辑。

Sequential 概括下来就是“快、干净、易懂”。

3.4 自定义模块：forward 写逻辑，init 定义积木

只要模型逻辑稍微复杂，就应该写成自定义 Module。例如：线性层 → ReLU → Dropout → 线性层。

1
2
3
4
5
6
7
8
9
10
11
12
13
class FeedForward(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim, p=0.1):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, out_dim)
        self.dropout = nn.Dropout(p)

    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

模块内部的分工： __init__声明积木（哪个层、哪些参数）；forward()说明积木如何拼接使用。一个常见错误示例：

1
2
3
def forward(self, x):
    fc = nn.Linear(128, 64)  # ❌ forward 里创建层
    return fc(x)

这样做的问题，每次 forward 都重新创建参数，无法被优化器管理，计算图混乱，无法保存模型，会导致模型完全不可训练。

规则：所有可训练的层必须放在 init 中定义，不要在 forward 中创建新层。

3.5 残差模块（Residual Block）：模块化设计的典型案例

残差模块（ResNet block）是一种有分支、有跳连的结构，极其适合用 nn.Module 表达。下面是经典 BasicBlock 的一个简化实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)

    def forward(self, x):
        identity = x                # 残差连接
        out = self.conv1(x)
        out = self.relu(out)
        out = self.conv2(out)
        out = out + identity        # 残差相加
        out = self.relu(out)
        return out

结构示意图用 mermaid 表达如下：

graph LR
    A[input] --> B[Conv2d]
    B --> C[ReLU]
    C --> D[Conv2d]
    A -- 跳连 --> D
    D --> E[ReLU]
图 3-2 残差结构（Residual Block）示意图

左侧输入 x一路向右进入 Conv → ReLU → Conv，输入 x 同时走一条上方“直连”路线到第二个 Conv 的输出。两条路径在 “Add” 节点相加，最后再过 ReLU。这个例子体现了模块化带来的最大价值：forward 里可以写任意计算图，而不会被 Sequential 限制。

3.6 子模块与参数管理：为什么优化器不需要你手动列参数？

任何写成 nn.Module 的模型，所有 nn.Parameter 会被自动注册，且所有子模块的参数会自动被包含。model.parameters() 会递归遍历整个 Module Tree。我们不需要手动把参数列出来：

1
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

哪怕模型结构是这样：

1
2
3
4
5
6
7
8
Model
 ├── Backbone
 │    ├── Block1
 │    └── Block2
 │          ├── conv1
 │          └── conv2
 └── Head
      └── fc

优化器都会自动找到所有 conv 层的权重以及所有线性层的权重，还有所有可训练参数。其实我们不需要也不应该去手动整理参数列表。

3.7 使用 register_buffer：存储非可训练张量

有些张量需要随模型保存，但不需要梯度，比如BatchNorm 的 running mean / var、位置编码（固定）、掩码矩阵（mask）、统计量，或一些常量表。这类张量应该使用：

1
self.register_buffer("pos_encoding", pe_tensor)

它们会随着模型保存、加载，不会出现在 model.parameters() 中，也不会被优化器更新，会随着 .to(device) 自动迁移至目标设备。这是很多人容易忽略，但非常重要的模块化工具。

3.8 模块化带来的可组合性：搭建大型模型的方式

我们可以通过 Module + Module + Module 的方式构建复杂结构：

1
2
3
4
5
6
7
8
9
10
11
12
class TransformerEncoderLayer(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.attn = SelfAttention(dim)
        self.ffn = FeedForward(dim, 4*dim, dim)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

再用多个 encoder layer 堆叠：

1
2
3
layers = nn.ModuleList([
    TransformerEncoderLayer(dim) for _ in range(6)
])

这展示了模块化核心价值：复杂网络 = 小模块的组合，而不是一个巨大函数。

小结：nn.Module 是 PyTorch 的“积木体系”
本章从多个角度说明了为什么 PyTorch 的模型必须建立在 nn.Module 之上。首先，一个模型在本质上是一棵由子模块组成的树状结构，模块之间可以层层嵌套、分层组织，使参数管理、结构表达以及模型组合都变得清晰而有序。在这种体系中，init 用于声明模型所需的“积木”，而 forward 则负责定义计算的实际流程；也正因为职责清晰，任何可训练参数都不应在 forward 中被临时创建。

对于结构简单、按顺序堆叠的网络，nn.Sequential 提供了最直观的表达方式；而更复杂的模型则依赖自定义 Module 来实现，如具有分支结构、跳跃连接的残差模块，它们展示了模块化设计在表达复杂拓扑时的巨大威力。得益于这一体系，模块能够自动记录并管理所有参数，使保存、加载、分布式同步以及优化器的构建全部变得自动而可靠。此外，在某些情况下需要保存状态但又不希望它参与训练，这时 register_buffer 提供了一种专门存放“非可训练张量”的方式，保持模型结构与参数管理的统一性。

4. 高性能数据管线：让 GPU 不再空等数据

在前面章节中，我们一直专注于“模型如何计算”，但深度学习训练中另一个关键问题是：GPU 是否真的在“全速计算”？还是大部分时间都在“等待数据”？ 现代 GPU 的算力远超 CPU，当数据预处理、加载、传输不够快时，GPU 会出现大量空闲。下面是一些常见的问题：

  • GPU 利用率忽高忽低；
  • 每个 batch 之间训练卡顿；
  • DataLoader 速度远远弱于模型速度；
  • num_workers 设成默认 0，导致 CPU 串行加载数据。

本章目标是回答以下几个问题

  • Dataset 和 DataLoader 在 PyTorch 中扮演什么角色？
  • 如何正确利用多进程、预取、Pinned Memory？
  • 如何构建可扩展、高性能的数据加载体系？
  • 如何验证数据管线的吞吐是否足够快？

4.1 PyTorch 的数据管线：Dataset ≠ DataLoader

PyTorch 将“数据的逻辑组织”和“数据的高效加载”彻底解耦：

  • Dataset：定义“给定索引 idx → 返回样本”
  • DataLoader：控制“如何按批量、多进程、异步预取地取数据”

这两者的解耦是提升性能的关键。

4.2 Dataset：定义数据的逻辑结构

一个 Dataset 最重要的是三个信息：

  1. 数据有多少条？ → len
  2. 如何通过索引访问一条样本？ → getitem
  3. 是否需要 transform包括数据要如何预处理？

例如加载 CIFAR-10 数据集的 Dataset：

1
2
3
4
5
6
7
8
9
from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
])

train_set = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)

4.3 DataLoader：数据的物理加载方式

DataLoader 决定了：

  • 怎样把 Dataset 打包成一个个 batch；
  • 是否多进程并发读取；
  • 是否预先加载下一批数据（prefetch）；
  • 是否使用 Pinned Memory 加速 CPU→GPU 拷贝；
  • 是否 shuffle / drop_last。

其最基本用法如下：

1
2
3
from torch.utils.data import DataLoader

train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=4, pin_memory=True)

4.4 DataLoader 的内部工作机制

DataLoader 的性能主要取决于三个因素：

  • num_workers：多进程并发取样本
    默认 num_workers=0，意味着所有数据加载 = 主线程串行执行。如果 transform 很复杂或者图片很多，这会极其慢。在设置为多进程后，每个 worker 进程都会调用 Dataset.getitem()。比如 num_workers=4，这时4 个进程会同时读取数据，然后由主进程负责拼 batch、调度。
  • Prefetch（异步预取）：隐藏加载延迟
    当 GPU 正在训练当前 batch 时，DataLoader 的 worker 会预先加载下一批数据。如果没有预取：训练结束 → 等待加载 batch → 再训练 → 再等待……。有了预取以后，训练 batch N 的同时，worker 正在准备 batch N+1。这样可以极大减少训练间隙。
  • pin_memory：加速 CPU→GPU 传输
    如果模型在 GPU 上训练：pin_memory=True。可以显著减少以下过程的开销：batch.cpu_tensor.to("cuda")。Pinned Memory 是一种“页锁定内存”，GPU 能用 DMA 进行零拷贝传输，比普通内存快很多。

4.5 num_workers 差异的实验

下面我们用 CIFAR-10 + 复杂数据增强 + 大 batch，让 CPU 负载足够高，从而体现并行加载加速效果。步骤：

  1. 使用 RandomCrop + RandomHorizontalFlip（CIFAR 常用增强）
  2. batch_size 设置为较大值（如 512）
  3. 迭代完整数据集（一个 epoch）
  4. 测量耗时
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import time
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

def benchmark(num_workers):
    transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor()
    ])

    dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)

	# 大 batch，让 CPU 压力更大
    loader = DataLoader(dataset, batch_size=512, shuffle=True, num_workers=num_workers, pin_memory=True)

    start = time.time()
    for _ in loader:
        pass  # 只衡量数据加载时间
    return time.time() - start

for w in [0, 2, 4, 8]:
    t = benchmark(w)
    print(f"num_workers={w}: {t:.2f}s")

日志

1
2
3
4
num_workers=0: 12.5s
num_workers=2: 6.9s
num_workers=4: 4.8s
num_workers=8: 4.5s

可以看到num_workers=0 明显最慢，增加到 4 有明显提升，8 稍有收益，但增加过多会引起调度开销。当 transform 复杂、batch 大时，多进程还可以成倍提升数据吞吐。 下面是关于选择 num_workers 的一些经验：

数据规模 推荐设置
小数据（MNIST, Fashion-MNIST） 0 或 2 即可
中等数据（CIFAR, CelebA） 4–8
大数据（ImageNet） 8–16
自定义大图 + 重数据增强 num_workers ≈ CPU 核心数

当然最终还是以实际 benchmark 为准。

4.6 高性能数据加载的最佳实践 Checklist

以下设置适用于大部分 GPU 训练场景：

  • 设置 num_workers > 0
  • 开启 pin_memory=True
  • 大 batch 情况下开启预取
  • 在 Windows/Mac 使用 if name == “main“: 包裹入口
  • 在 GPU 上训练时，不要让 transform 过度复杂（可用 GPU 加速增强库）
  • 多 GPU 训练时可以使用 DistributedSampler
  • 少用 Python 原生操作（如 for 循环），尽量用 torchvision transforms

小结：数据管线是训练性能的关键瓶颈
本章通过 Dataset、DataLoader、多进程、预取以及 Pinned Memory，说明了深度学习训练的性能瓶颈常常不在 GPU，而在数据加载。 现在我们应该清楚：

  • Dataset 只是描述数据是什么；
  • DataLoader 决定数据如何被高效地“喂”给 GPU；
  • num_workers 与 pin_memory 是性能优化的核心；
  • CIFAR-10 实验展示了并行加载的实际提升；
  • DataLoader 本质上是一个“异步生产者-消费者系统”。

5. 动态计算图与自动微分：PyTorch 的灵魂机制

在第 1 章我们跑通了训练流程，在第 2 章理解了 Tensor 与梯度，在第 3–4 章学习了模型和数据管线。现在我们继续PyTorch： 为什么 loss.backward() 就能让整个模型自动求梯度？ 以及：

  • 计算图是如何构建的？
  • grad_fn 是什么？
  • backward 是沿着什么路径传播？
  • detach() 为什么能“切断”梯度？
  • in-place 操作为什么危险？
  • 为什么 PyTorch 能支持动态控制流（if/for）？

这些构成了 PyTorch 相比静态框架最大优势：Define-by-Run —— 执行即构图（动态图）。 本章将把自动微分的运行机制彻底讲清楚。

5.1 Define-by-Run：PyTorch 的动态图思想

我们在 Python 中写什么代码，PyTorch 就记录什么计算路径。比如：

1
2
x = torch.tensor(3.0, requires_grad=True)
y = x * 2 + 1

这两行代码做了两件事：在执行数学计算的同时，构建一张“记录这次计算过程”的计算图（Graph）。后续调用y.backward()，这时PyTorch 就会从 y 的 grad_fn 出发，沿着计算路径反向遍历自动应用链式法则最终把 ∂y/∂x 放到 x.grad 中。

5.2 Tensor.grad_fn：每个张量都知道“自己来自哪里”

在 PyTorch 中，只要一个 Tensor 是 通过运算产生的（不是手动创建的），并且它参与了自动求导，那么它都会带上一个属性：grad_fn。举个例子：

1
2
3
4
5
6
x = torch.tensor(2.0, requires_grad=True)
y = (x + 3) ** 2

print(y.grad_fn)  # <PowBackward0>

print(y.grad_fn.next_functions) # ((AddBackward0, 0),)

y 是由平方操作（Pow）产生的PowBackward0，平方的输入来自加法操作（Add） → AddBackward0，加法的输入是叶子张量 x，x就没有 grad_fn。可以把 grad_fn 理解成：“这个张量是由哪个运算算出来的？backward 时应该沿着哪条路径回传？” 它是 反向传播的路线图。

在自动微分中，PyTorch 会把每一步前向运算都登记成一个“节点”，并将它们按计算顺序串在一起。执行 y.backward() 时，系统就会从 y 的 grad_fn 开始，一层层顺着 next_functions 倒推回去，自动完成链式求导。如果一个张量是直接创建的，如：

1
2
z = torch.tensor(5.0)
print(z.grad_fn)  # None

它就没有 grad_fn，因为它不是由计算生成的，是计算图的“叶子”，是反向传播最终停下来的节点。grad_fn 让 PyTorch 在 backward 时知道“这个张量从哪里来”，并据此完成自动求导。 这正是动态计算图能够“边执行边记录”的关键所在。

5.3 backward：沿计算图反向求导

理解 backward 的关键是：它并不是“凭空”算梯度，而是顺着 grad_fn 记录的计算路径，一步步把梯度从输出传回输入。继续使用同一个简单例子：

1
2
3
4
x = torch.tensor(2.0, requires_grad=True)
y = (x + 3) ** 2
y.backward()
print(x.grad)  # tensor(10.)

数学上：
$$y = (x + 3)^2,\quad \frac{dy}{dx} = 2(x+3)=10$$
但 PyTorch 的求导方式不是先推公式再代入数值，而是沿着计算图“倒着走”：

  1. y 的 grad_fn 是 PowBackward0，说明 y 来自一次“平方”
  2. backward 会调用平方的反向规则，把梯度传给它的输入 (x+3)
  3. 这个输入由 AddBackward0 生成，所以继续向前一步
  4. backward 再调用加法的反向规则，把梯度传给最终的叶子张量 x
  5. 得到的 dy/dx 就被写入 x.grad

换句话说，backward 做的工作就是：根据 grad_fn，把梯度一层一层沿原路传回去。为了更直观，可以画成这样一条路径：

flowchart LR
    %% --- 样式定义 ---
    classDef forward stroke:#4A90E2,stroke-width:2px,color:#111
    classDef backward stroke:#F5A623,stroke-width:2px,stroke-dasharray:5 5,color:#111
    classDef tensor fill:#fff,stroke:#666,stroke-width:1.5px,rx:6,ry:6

    %% --- Forward Path ---
    X[x]:::tensor -->|Add| A[(x + 3)]:::tensor
    A -->|Pow| Y[y]:::tensor

    %% --- Backward Path (虚线) ---
    Y-.->|PowBackward0| A
    A-.->|AddBackward0| X
图 5-1 forward/backward 路径

forward 是从 x 走到 y，backward 则沿着 grad_fn 反方向从 y 走回 x。每经过一个节点，就调用对应的“反向算子”（如 PowBackward0、AddBackward0）来计算局部梯度。所以最后 x.grad 得到的 10，既不是手算的，也不是 PyTorch 硬编码的，而是自动沿图反传的结果。

一句话总结：backward = 顺着 grad_fn 标记的计算图，把梯度按链式法则自动传回输入张量。 这样自动微分才得以实现，不需要你为每个模型手写梯度公式。

5.4 动态计算图：支持 Python 控制流的关键

在静态图框架（早期 TensorFlow）中，if/for/while 都必须用特殊算子表达。而 PyTorch 动态图的核心优势是：控制流写在 Python 中即可自然成为计算图结构。 例如：

1
2
3
4
5
6
7
8
9
def dynamic_net(x):
    y = x * 2
    for i in range(int(x.item())):
        y = y + i
    return y

x = torch.tensor(3.0, requires_grad=True)
y = dynamic_net(x)
y.backward()

这里计算图结构取决于x 的值，以及for 循环次数和分支路径。每一次 forward 都会生成一张全新的计算图。图示如下：

flowchart LR
    X[x] --> A[*2]
    A --> B[+0]
    B --> C[+1]
    C --> D[+2]
图 5-3 动态控制流下实时生成的计算图

x=3 → for 循环执行 3 次，图结构动态展开为 “+0 → +1 → +2”，若 x=1，就会只有一次 “+0”，forward 决定图结构，图是“实时搭建”的。

5.5 detach()：切断梯度传播

detach() 的作用是：保留张量的数值，但把它从当前计算图中“摘”出来，让它不再参与梯度回传。 例如，有些操作需要用到参数的“数值”，但并不希望更新它的梯度，最典型的场景是 EMA（指数滑动平均）：

1
shadow = decay * shadow + (1 - decay) * param.detach()

又或者我们只是想把中间特征拿出来可视化，而不是让它影响反向传播：

1
2
with torch.no_grad():
    y = model(x)

这些操作背后都是同一件事：使用它的值，但不要它的梯度。 我们来看一个更直观的例子：

1
2
3
x = torch.tensor(3.0, requires_grad=True)
y = x * 2
z = y.detach()  # z 与计算图脱离

这里y 仍然属于从 x 出发的计算图，因此 backward 会从 y 回到 x，z 只是“拿到了 y 的值”，但梯度链路被切断，z.grad_fn = None。可以用下面的图来理解（虚线表示“没有梯度关系”）：

flowchart LR
    X[x] -->|*2| Y[y = x*2]
    Y -.->|detach| Z[z 无梯度]
图 5-4 detach() 生成不带梯度的新张量

y 与 x 的连接是正常计算图，它们会一起参与 backward，z 虽然数值等于 y，但已完全脱离计算图，在反向传播中不会影响任何梯度。

5.6 requires_grad_()：动态开启/关闭梯度

冻结模型部分参数：

1
2
for p in model.backbone.parameters():
    p.requires_grad_(False)

再开启：

1
p.requires_grad_(True)

适用于冻结预训练 backbone，分阶段训练，或者手动控制梯度流动

5.7 autograd.set_detect_anomaly()：调试梯度的利器

当反向传播中出现NaN、Inf，或者某个算子 backward 报错。使用：

1
2
torch.autograd.set_detect_anomaly(True)
loss.backward()

PyTorch 会提示我们：哪个算子出现了问题，该算子的 forward/backward 在哪一行，可能的原因是什么。这是训练大模型时必备的调试工具。

5.8 实战：从梯度流判断模型是否“正常学习”

一个模型是否在学习，不看 loss，也不看 acc，而看：梯度是否健康地沿网络流动。 简单监控梯度范数：

1
2
3
4
5
6
total_norm = 0
for p in model.parameters():
    if p.grad is not None:
        total_norm += p.grad.data.norm().item()

print("Grad Norm:", total_norm)

异常信号：

  • 梯度范数接近 0 → 梯度消失
  • 梯度范数巨大（>1000） → 梯度爆炸
  • 某层梯度一直为 0 → 计算图断裂 / in-place 破坏
  • 突然出现 NaN → 数值稳定性问题（学习率太大等）

小结：PyTorch 的灵魂是动态图
本章系统阐述了 PyTorch 自动微分机制的运行方式，核心思想在于它采用了“Define-by-Run”的动态图范式，即计算图不是预先定义好的，而是在每一次 forward 中即时生成。前向传播结束后，会得到一张完整且只属于当前计算的一次性计算图；在 backward 执行完毕后，这张图便会被立即释放，从而保持 PyTorch 的灵活性和高效性。

在这张图中，每个 Tensor 都带有自己的来源信息，grad_fn 记录了它是由哪一步运算生成的。反向传播会根据这条链条，从输出位置沿着计算图自动向后追踪，依次计算梯度并传递至所有参与运算的张量。由于图是“运行时构建”的，因此 PyTorch 天然支持 Python 语言的全部控制流，包括循环、条件分支、动态结构等，这是静态图框架难以做到的。

本章也介绍了控制计算图行为的关键机制，例如 detach 可以切断梯度路径，让某个张量不再参与反向传播；requires_grad 决定了哪些张量需要追踪梯度；而 autograd 的异常检测工具则帮助我们在求导出错时快速定位问题。

在调试层面，监控梯度流是理解模型是否真正“在学习”的重要手段。无论是梯度消失、梯度爆炸，还是某些层完全没有梯度更新，都可以通过仔细检查梯度流动情况来定位和解决。

6. PyTorch 调试：从梯度到数值稳定

深度学习代码能跑并不意味着能“正常训练”。真正棘手的问题往往出现在：

  • 模型不收敛
  • 梯度莫名其妙为 0 或爆炸
  • loss 出现 NaN/Inf
  • 某层完全不更新参数
  • backward 报异常，看不懂 traceback
  • 修改一点代码后训练结果完全不同
  • DataLoader 或模型结构导致图断裂

训练真实模型时，这些问题比单纯写模型碰到的问题还常见。下面讲讲 PyTorch 的调试方法，来帮助定位“训练不稳定”的真正原因。

6.1 调试核心：训练是一个信号流动过程

训练的本质是：前向是数据的流动，反向是梯度的流动。 模型能否学习，核心是要看梯度是否能顺畅地流通。因此，调试的核心不是“loss 多少”，而是：

  • 梯度是否在每一层都有合理的值？
  • 数据是否有异常（极值 / 常数 / 错位）？
  • 哪一层截断了梯度？
  • 是否存在不稳定数值（Inf / NaN）？
  • 模型参数是否被冻结？

我们可以用一个简化示意图来理解训练信号流：

flowchart LR
    Input --> F1[Layer 1]
    F1 --> F2[Layer 2]
    F2 --> F3[Layer 3]
    F3 --> Loss

    Loss -.-> B3[Grad 3] -.-> F3
    B3 -.-> B2[Grad 2] -.-> F2
    B2 -.-> B1[Grad 1] -.-> F1
    B1 -.-> Input
图 6-1 训练的“信号流”模型：前向数据流（实线）与反向梯度流（虚线）

从左到右是前向数据流 (forward)，从右到左是反向梯度流 (backward)。如果某一处发生断裂，就会影响整条链路。整个训练是前向 + 反向 的循环

6.2 第一类问题：梯度异常（0、巨大、NaN）

最常见、也最致命的问题来自“梯度异常”。我们可以通过一个通用监控工具来检测：监控梯度范数（Gradient Norm）

1
2
3
4
5
6
7
8
9
def grad_norm(model):
    total = 0
    for p in model.parameters():
        if p.grad is not None:
            total += p.grad.data.norm().item()
    return total

# 在训练循环中调用：
print("Grad Norm:", grad_norm(model))

如何判断梯度异常？

异常类型 症状 原因
梯度为 0 模型不更新、loss 不下降 激活函数饱和、图断裂、忘记 requires_grad
梯度 爆炸 loss 变成 Inf / NaN 学习率太大、指数增长模块、RNN 不稳定
梯度中有 NaN loss=NaN / backward 崩溃 0 除、log(0)、无效操作、数值上溢

特别是如果梯度突然爆炸，基本一定是数值不稳定，梯度一直是 0 → 90% 是计算图断了。

6.3 第二类问题：loss = NaN / Inf（数值不稳定）

AI 工程中最常见的问题：

1
Epoch 12: loss = nan

数值不稳定通常来自触发无效数学运算

操作 问题点 示例
log(0) 结果为 -inf F.log_softmax 一般会做保护
除以 0 inf x / eps 要写成 +1e-8
指数 e^x 很容易溢出 attention score
平方和累加 大数堆叠 ((x**2).sum())**2
  • Softmax/Logits 数值溢出
1
2
3
4
# 错误写法
prob = torch.exp(x) / torch.exp(x).sum()
# 正确写法
prob = F.softmax(x, dim=-1)
  • 没有使用 CrossEntropyLoss
1
2
3
4
5
6
# 不这么写
prob = F.softmax(logits)
loss = -(y * prob.log()).sum()
# 一般用CrossEntropyLoss做数值稳定处理
criterion = nn.CrossEntropyLoss()
loss = criterion(logits, target)

6.4 第三类问题：模型参数“没有更新”

常见原因如下：

  • 忘记设置 requires_grad
1
2
3
4
# 默认不追踪梯度
x = torch.randn(10) 
# 应该是这样
x = torch.randn(10, requires_grad=True)
  • 在 forward 中创建层
    这样会导致Optimizer 无法管理以及参数不更新。
1
2
3
def forward(self, x):
    fc = nn.Linear(128, 64)  # 每次 forward 又创建新参数
    return fc(x)
  • detach() 用错位置
1
x = self.encoder(x).detach()  # 移除了梯度
  • in-place 破坏计算图
1
x.relu_()  # 原地覆盖

6.5 第四类问题：计算图断裂（Graph Break）

图断裂会导致梯度不再回传，症状是某几层梯度一直为0，或者loss 不下降。常见原因有以下几点：

  • Tensor 转 numpy
1
2
3
4
# 彻底断裂，梯度丢失
x = x.numpy()   
# 明确表示“不需要梯度”
x = x.detach().cpu().numpy()
  • Tensor.item() 参与运算
    float 不在计算图中，梯度链路断裂。
1
y = x * x.item()   # item() 返回 Python float
  • detach()
1
2
y = x.detach()
z = y * 2        # 不会回传给 x

6.6 第五类问题：backward 报错（定位算子）

PyTorch 提供的调试工具：

  • autograd.set_detect_anomaly(True)
    当 backward 出错时它会打印出 forward 中的问题算子，显示具体行号，说明哪个节点梯度计算失败。
1
2
torch.autograd.set_detect_anomaly(True)
loss.backward()

6.7 使用 Hook 调试：捕获任意层的输入/输出/梯度

Hook 是 PyTorch 中一个低调但极其强大的功能。

  • 注册梯度 hook
    可用于查看，某层梯度是否消失，数值是否溢出，或者是否为 NaN
1
2
3
4
def print_grad(grad):
    print("Grad:", grad.norm())

x.register_hook(print_grad)
  • 注册 forward hook
    一般用于检测：是否某层输出全部是 0、是否某层输出巨大（引发爆炸）、是否某层被误触发 in-place 改写
1
2
3
4
def fwd_hook(module, inp, out):
    print(module.__class__.__name__, out.mean().item())

layer.register_forward_hook(fwd_hook)

6.8 数据问题：最容易被忽略的训练失败原因

常见数据异常：

  • 数据全是 0/255
    图像输入未归一化 → 模型无法训练：
1
transforms.Normalize(mean, std)
  • 标签错位
    train_loader 输出 (data, label)
1
2
for label, data in train_loader:
   …
  • Batch Size 太小
    BN（BatchNorm）在 batch_size 太小时会不稳定，导致发散。
  • shuffle=False
    模型学习不到统计多样性，可能出现奇怪收敛轨迹。

小结：PyTorch 调试的系统方法
通过本章内容，我们建立了一套系统化的 PyTorch 调试思维。首先，从梯度出发，可以通过观察 grad_norm 来判断整体梯度是否健康，也可以借助 hook 精确查看某一层的梯度流动情况。在此基础上，需要特别警惕 requires_grad、detach 的误用，以及任何可能导致计算图断裂的操作，这是调试梯度问题的核心。

在数值稳定性方面，需要避免出现 log(0)、除以零、指数运算溢出等典型错误，并确保 softmax 与 crossentropy 的使用方式正确，同时对 exp、平方、范数累积等操作保持谨慎，以防止出现数值爆炸。

当模型出现“参数不更新”的情况时，应从结构与逻辑两方面排查：例如是否在 forward 中意外创建了新的层，是否错误冻结了参数，是否因为 in-place 操作覆盖了中间值，或是否因 detach 阶段性切断了梯度。这些问题往往直接影响模型能否成功学习。

如果 backward 阶段出现异常，则可以借助 anomaly 工具捕获具体的算子错误，也可以通过逐层 hook 逐步缩小排查范围，最终找到导致求导失败的节点。

调试过程中还必须关注数据本身是否正常，包括输入的归一化是否正确、标签是否错位、batch 是否过小导致统计不稳定，以及 shuffle 是否在训练阶段被正确启用。数据问题往往是训练异常的根源之一。

7. 性能：从 Eager 执行到编译加速

PyTorch 之所以深受研究者和工程师喜爱，来源于其“动态计算图 + Pythonic API”带来的自由度。虽然Eager（即时执行）模式容易调试，却缺乏静态结构，难以进行深度优化。

PyTorch 2.0 之后，性能有了一个提升：在保持 Eager 的灵活性，同时获得接近静态图框架的性能。 下面我们将从 Eager → TorchScript → torch.compile → AMP → Profiling为主线，来理解PyTorch究竟是如何提升性能的。

7.1 训练是如何执行的？Eager 模式的本质

Eager Execution（即时执行）是 PyTorch 的核心，在 Eager 中，这段代码每一步都立即执行：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# @：matmul 调用
# +：add 调用
# relu：activation 调用
y = torch.relu(x @ w + b)
```  
每个算子都需要 Python 调度 -> C++ -> CUDA kernel 上下文切换。这样虽然灵活，但有两个天然瓶颈：
1. 无法提前优化计算图，算子已经在执行，没有静态结构可以整体优化。
2. Python 调度开销大，每一个算子都是一次独立 Python 调用。

所以 Eager 的优点是：极易调试、控制流友好、与 Python 一致。缺点是速度不如静态图、算子融合有限、Kernel 数量多还会带来串行 overhead，**这就是“灵活 vs 性能”的根本矛盾。**
### 7.2 TorchScript：第一次尝试让 PyTorch 拥有静态图
在 PyTorch 1.x 时代，官方曾尝试通过 TorchScript 提供静态图式加速。TorchScript 有两种方式：
- Tracing（跟踪法），优点是简单但缺点是无法处理控制流（if/for）。
```python
	traced = torch.jit.trace(model, example_input)
  • Scripting（脚本化），优点是支持控制流，但缺点语法限制多、理解难度大、难调试。
1
scripted = torch.jit.script(model)

TorchScript 最大的问题是： 用户必须迁就框架，而不是框架适应用户。 因此 TorchScript 更多用于“部署”，而非“提升训练性能”。

7.3 PyTorch 2.0：torch.compile 的出现

从 PyTorch 2.0 开始，性能有了一个很大的提升：不改模型代码，通过编译自动获得加速。 启用方式很简单：

1
model = torch.compile(model)

这样就可以直接训练，不需要改变任何 forward 或训练脚本。这是 PyTorch 性能策略的终极目标：

  • 让用户写完全“Eager 风格”的代码
  • PyTorch 在背后捕获计算图并优化
  • 用户获得 2×–3× 实测加速

7.4 torch.compile 的三段式架构

PyTorch 2.0 的编译体系由三个关键模块组成：

  1. TorchDynamo：捕获 Python 层下面的计算图
    Dynamo 会拦截 Python 执行，将每次 tensor 运算捕获为 FX Graph（中间表示）。它能处理：

    • if / for / while 等控制流
    • 动态形状（dynamic shapes）
    • Python side effect

    这是 TorchScript 无法做到的。

  2. AOTAutograd：提前构建反向图
    它会生成前向和反向的静态图，便于做算子融合与优化。相当于：

1
forward_graph, backward_graph = AOTAutograd(fx_graph)

  1. Inductor：生成高性能 CUDA/CPU Kernel。Inductor 会：

    • fuse（融合）多个算子
    • 减少 kernel launch 数量
    • 用 Triton 自动生成 CUDA Kernel
    • 减少 memory access
    • 提升吞吐率

    这部分性能提升非常可观。

7.5 混合精度训练（AMP）：Tensor Core + 更少显存 + 更高速

混合精度训练（Automatic Mixed Precision）已经成为现代训练的“默认配置”。PyTorch 提供两类 AMP：

  • FP16 AMP
  • BF16 AMP（更稳定、推荐使用 A100/H100）
1
2
3
4
5
6
7
8
9
10
11
12
13
14
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for x, y in loader:
    optimizer.zero_grad()

    with autocast():
        pred = model(x)
        loss = criterion(pred, y)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

AMP 的优势：

  • 使用 Tensor Core → 更高吞吐
  • FP16 权重减少显存占用
  • 速度提升 1.3×–2×
  • batch size 增大
  • 大模型训练更稳定

结合 torch.compile → 速度会进一步提升。

7.6 Gradient Accumulation：显存不足时的“大 batch 技巧”

如果显存不够训练大 batch，可以用梯度累积：

1
2
3
4
5
6
7
8
9
10
11
12
accum_steps = 4

for i, (x, y) in enumerate(loader):
    with autocast():
        loss = model(x, y) / accum_steps

    scaler.scale(loss).backward()

    if (i + 1) % accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

你用 4 次 batch size=32 的 forward，就能模拟：

1
effective batch = 32 × 4 = 128

这在小显存设备（如 8GB GPU）上非常实用。

7.7 DataLoader 性能：减少 Python overhead，提升吞吐

DataLoader 性能影响整个训练速度的 20%–50%。最关键参数：

  • num_workers：数据加载并行度
  • pin_memory=True：加速 H2D 传输
  • prefetch_factor：更高的预加载吞吐
  • persistent_workers=True：减少 worker 重建开销

推荐配置：

1
DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True, persistent_workers=True)

7.8 Profiling：性能优化不再靠猜测

PyTorch Profiler 能详细显示：

  • 哪些算子最慢
  • DataLoader 是否成为瓶颈
  • CUDA kernel 启动频率
  • CPU/GPU 并行度
1
2
3
4
5
6
7
8
9
import torch.profiler as profiler

with profiler.profile(
    activities=[profiler.ProfilerActivity.CPU, profiler.ProfilerActivity.CUDA],
    record_shapes=True
) as prof:
    output = model(x)

print(prof.key_averages().table(sort_by="cuda_time_total"))

所有性能优化都应该先 profile，再决策。

7.9 性能优化 Checklist

在模型加速方面，可以优先启用 torch.compile 来减少 Python 调度开销，并使用自动混合精度（AMP）获得显著的吞吐提升。同时，应尽量减少 Python 层面的条件分支，避免反复创建新的 Tensor，以减少构图与调度的额外成本。此外，除非必要，不应依赖 inplace 操作，以免破坏 autograd 或导致不必要的重新构图。

在数据吞吐方面，DataLoader 的并行加载至关重要。合理设置 num_workers 可以明显提升数据准备速度，pin_memory=True 则有助于加快 CPU 到 GPU 的数据传输。如果显存充足，可以采用更大的 batch；显存有限时，则可以通过梯度累积来模拟大批量训练。同时，分布式训练（DDP）能让吞吐进一步扩大，适用于多卡甚至多机环境。

在显存优化上，AMP 能减少 30%–50% 的显存压力，是最直接有效的手段。对于更大模型，可以使用 Gradient Checkpointing 以计算换显存；必要时还可以从 Adam 切换为占用更低的 SGD。此外，清理不必要的缓存、减少多余的中间激活，也有助于提升整体可训练规模。

在数值稳定性方面，GradScaler 是混合精度训练中保证稳定性的关键工具，应默认开启。对于 softmax 或 logits 等关键算子，尽量避免直接使用 FP16，以降低数值溢出风险。如果硬件支持（如 A100/H100），BF16 通常是兼具性能和稳定性的更佳选择。

小结：PyTorch 性能进入“编译时代”
本章构建了一套系统化的性能优化思路。我们首先理解了 Eager 模式的优势与瓶颈：它极其灵活，但在大规模训练中难以避免 Python 调度带来的开销。随后回顾了 TorchScript 作为早期静态图方案的局限性，也正因为这些限制，PyTorch 2.0 才通过 torch.compile 真正实现了“动态图与高性能”兼得的目标，让模型在保持原有写法的前提下获得编译器级别的加速。

在此基础上，我们进一步引入 AMP，使训练在吞吐、显存和数值稳定性之间获得最佳平衡；同时讨论了梯度累积与 Gradient Checkpointing，在大模型训练中，它们是突破显存瓶颈的关键技术。数据加载侧的优化同样重要，通过合理的 DataLoader 设置，我们能够解决输入阶段的性能瓶颈，让 GPU 不再等待数据。

在高性能训练中，Profiler 是定位瓶颈不可或缺的工具，它让优化不再依靠猜测，而能基于可视化证据进行决策。最后，通过一份性能 Checklist，我们将所有策略系统化、可执行化，使训练性能能够被清晰管理。

总结一句：PyTorch 的性能优化，正在从手工技巧转向编译器驱动的系统智能。

8. 工程生态：从实验到部署的优雅闭环

深度学习项目远不止“训练一个模型”那么简单。在真实工程中，一个模型往往需要经历：

1
训练 → 评估 → 记录 → 管理 → 导出 → 部署 → 监控

这是一条“完整生命周期（ML Lifecycle）”，PyTorch 的生态正是围绕它构建的。

本章目标是构建训练端与工程端的桥梁，让 PyTorch 代码真正进入“可复现、可共享、可部署、可维护”的工程状态。

8.1 模型保存与加载：从 state_dict 开始的工程基石

PyTorch 的核心设计思想之一：模型 = 代码结构（类） + 权重（state_dict） 。这比 TensorFlow 的“冻结图”更加灵活。

8.1.1 state_dict：官方推荐、最安全、最灵活的保存方式

保存：

1
torch.save(model.state_dict(), "model.pth")

加载：

1
2
3
model = MyModel()
model.load_state_dict(torch.load("model.pth"))
model.eval()

这种方式的好处：

  • 只保存权重，不保存代码逻辑
  • 不依赖具体 Python 环境
  • 不受 pickle 安全问题影响
  • 易于迁移学习、fine-tune
  • 文件体积更小

很多初学者犯的错误是：

1
torch.save(model, "model.pth")  # 强依赖当前代码环境

这种方式不推荐，因为：

  • 需要模型类在加载环境中同名存在
  • pickle 安全风险
  • 环境变化时容易失效

工程实践：推荐保存 state_dict。

8.1.2 训练断点：保存优化器状态（必要时）

训练中断？不想重来？只需保存 optimizer 状态：

1
2
3
4
5
torch.save({
    "epoch": epoch,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict()
}, "ckpt.pth")

加载：

1
2
3
4
ckpt = torch.load("ckpt.pth")
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
start_epoch = ckpt["epoch"] + 1

适用于：

  • 大模型训练
  • 昂贵训练任务
  • 云上训练（spot 实例随时中断）

8.2 PyTorch Lightning：让训练循环从“脚本”变成“系统”

原生 PyTorch 灵活，但工程端常见问题：

  • 训练循环重复代码多
  • 逻辑（模型）与工程（训练）混杂
  • hyper-parameter 到处漂移
  • 多 GPU 分布式写起来复杂
  • 日志、回调、checkpoint 杂乱

Lightning 的理念：把科学部分（模型）与工程部分（训练逻辑）分离。

8.2.1 LightningModule：模型定义更清晰

一个 Lightning 模型只负责：

  • forward
  • loss
  • backward（Lightning 管）
  • optimizer
  • metrics

示例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import pytorch_lightning as L

class LitMLP(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = SimpleMLP(784, 256, 10)
        self.loss_fn = nn.CrossEntropyLoss()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        pred = self(x.view(x.size(0), -1))
        loss = self.loss_fn(pred, y)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

训练只需：

1
2
trainer = L.Trainer(max_epochs=5, accelerator="gpu", devices=1)
trainer.fit(model, train_loader, val_loader)

8.2.2 Lightning 自动处理的工程事务

Lightning 可以替我们做：

  • 混合精度 AMP
  • GPU / 多机多卡
  • 梯度累积
  • 梯度裁剪
  • checkpoint 自动保存
  • 可重启训练
  • 回调机制
  • 日志集成（TensorBoard / W&B）
  • 训练 loop 管理

写的代码变少了，但工程质量更高、更专业。

8.3 Hydra：真正可复现的配置管理系统

深度学习项目常见问题：

  • 超参数散落在代码、命令行、config 文件
  • 复制一份 config 就乱套
  • 实验版本无法复现
  • 尝试不同参数组合时必须改 Python 代码

Hydra 的理念：让配置成为“一等公民”。

8.3.1 典型 Hydra 配置结构

configs/train.yaml：

1
2
3
4
5
6
7
model:
  hidden_dim: 256
  lr: 1e-3

data:
  batch_size: 64
  num_workers: 4

在 Python 中使用：

1
2
3
4
@hydra.main(config_path="configs", config_name="train")
def main(cfg):
    model = SimpleMLP(784, cfg.model.hidden_dim, 10)
    loader = build_dataloader(cfg.data.batch_size)

想修改配置？

1
python run.py model.hidden_dim=512 data.batch_size=128

不用改代码，不用复制 config，完全可复现。

8.4 实验追踪：让科研从“文件系统”进入“系统平台”

没有实验追踪的训练过程是不可控的。我们可能会遇到：

  • 哪次实验最好？
  • 上次的学习率是多少？
  • batch size 用的多少？
  • 最佳 checkpoint 在哪？
  • 团队成员怎么共享结果？

现代训练强烈建议使用：

  • Weights & Biases（W&B）
  • MLflow
  • TensorBoard（轻量）

8.4.1 W&B 最小示例

1
2
3
4
5
6
import wandb
wandb.init(project="mnist-demo")

for step, (x, y) in enumerate(train_loader):
    loss = train_step(x, y)
    wandb.log({"loss": loss})

Lightning 中更简单：

1
2
3
4
5
6
from pytorch_lightning.loggers import WandbLogger

logger = WandbLogger(project="mnist-demo")

trainer = L.Trainer(logger=logger)
trainer.fit(model, train_loader)

W&B 自动记录loss/acc 曲线、模型权重、配置参数、GPU/CPU 资源、模型对比、自动可视化。W&B 是目前最适合科研与中大规模工程团队的实验平台。

8.5 模型导出：从 PyTorch 到跨平台推理

训练好模型，只是旅程的 一半。接下来我们需要：

  • 在 web 推理？
  • 在 TensorRT 上优化？
  • 在手机端运行？
  • 在多语言环境（C++/Java）调用？
  • 与其他框架集成？

PyTorch 提供了两条路线：

8.5.1 ONNX：工业界最通用的深度学习交换格式

导出模型：

1
torch.onnx.export(model, example_input, "model.onnx")

ONNX 可用于：

  • 推理引擎（TensorRT、ONNX Runtime）
  • 移动端
  • C++ 推理
  • 云端服务

ONNX 会将模型转换成静态图格式，便于进一步优化。

8.5.2 TorchScript（部署侧）

尽管 TorchScript 在训练中热度降低，但在部署侧仍然非常有价值：

  • C++ LibTorch 推理
  • 移动端
  • 无 Python 的推理环境
1
2
scripted = torch.jit.script(model)
scripted.save("model.pt")

8.5.3 TensorRT：NVIDIA GPU 的终极加速方式

典型流程：

1
PyTorch model → ONNX → TensorRT Engine → 高性能 GPU 推理

TensorRT 能带来：

  • 2×–10× 推理加速
  • FP16 / INT8 量化
  • Kernel 融合

适用于大模型部署、推理服务。

8.6 TorchServe：生产级 PyTorch 推理服务

如果我们想将模型部署为 API 服务：

  • 高并发
  • 多模型多版本
  • 自动 batch
  • GPU 高效利用

那么官方给你的答案是：TorchServe（正式生产级）
目录结构：

1
2
3
4
5
model_store/
    mnist.mar      ← 打包模型
config/
    config.properties
serve.py

基本流程：

  1. 将模型打包为 .mar 文件
  2. 通过 torchserve –start 启动服务
  3. 通过 REST API 调用预测接口

TorchServe 是深度学习服务化的工业级方案。

工程闭环总结：PyTorch 的真正力量不只在训练
在这一整套工程流程中，我们已经从实验走向部署，构建起一个真正可落地的深度学习管线。PyTorch 的工程生态可分为几个关键层面：在训练阶段，state_dict 提供轻量而透明的权重管理方式，Lightning 让训练结构化并稳定可控，而 AMP、torch.compile 与 DDP 则构成了现代训练不可或缺的性能体系；在配置与复现方面，Hydra 负责统一且可组合的配置管理，W&B 与 MLflow 则提供系统化的实验追踪，使每一个实验都能被记录、对比与复现；在部署阶段，ONNX 实现跨平台兼容，TorchScript 支持 C++ 与移动端运行，TensorRT 在 GPU 上提供极致加速，而 TorchServe 则让模型真正运行在高并发的生产环境中。

一句话总结：PyTorch 并不是一个“训练框架”，而是一套完整的深度学习工程生态系统。

9. 实战整合：从模型训练到实验追踪的完整项目

在前 1–8 章中，我们从最基础的张量与自动微分开始，逐步了解了模块化的模型构建方式、高性能数据管线、动态计算图的调试能力、混合精度与 torch.compile 加速策略，并进入工程化生态，包括 Lightning、Hydra 和日志系统。

这些能力都为本章做铺垫——将所有知识整合成一个真正可用、可复现、可扩展、可部署的深度学习工程项目。

本章我们将学习如何构建一个专业级深度学习工程体系：

  • 清晰的项目结构
  • 完整的数据管线
  • 模块化模型定义
  • 工程化训练流程（Lightning + AMP + compile）
  • 实验可复现（Hydra 配置）
  • 自动保存最佳模型
  • 自动导出 TorchScript 部署模型（model.pt）
  • 本地推理脚本
  • 在线 API 推理（FastAPI）

这也是项目中的完整链路，最终产物可以直接实际应用。

9.1 为什么“完整项目”比“能训练模型”更重要？

许多初学者的训练代码往往是这样的，一个 800~2000 行的 train.py，模型、数据、训练循环、配置全部混合，改一点全盘崩溃，实验无法复现，无法部署，更别说团队协作。但在实际项目中，必须满足五个标准：

  1. 清晰结构（模块化） 模型、数据、训练、配置、部署解耦。
  2. 可复现（Hydra 配置 + Logging） 所有超参数有记录，实验可重现。
  3. 可扩展（更换模型/数据不破坏现有代码）
  4. 可加速（AMP、compile、DDP、多进程 DataLoader）
  5. 可部署（ckpt → TorchScript / ONNX → 服务）

9.2 项目结构设计

一个成熟的深度学习工程项目必须明确职责、形成模块化结构，本项目结构如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
project/
│── configs/                     # Hydra 配置  
│   └── config.yaml
│
│── src/
│   ├── data.py                  # 数据管线
│   ├── model.py                 # 模型定义
│   ├── lit_module.py            # LightningModule 
│   ├── train.py                 # 构建 Trainer + Checkpoint
│   ├── utils.py                 # 工具函数
│   └── __init__.py
│
│── deploy/
│   ├── export.py                # ckpt → model.pt / model.onnx
│   ├── inference.py             # 本地推理脚本
│   └── api.py                   # FastAPI 在线推理
│
│── run.py                       # 项目入口（训练 + 自动导出）
│── requirements.txt
│── README.md

这样架构就可以做到：

  • 修改模型 → 只改 model.py
  • 修改数据 → 只改 data.py
  • 修改训练参数 → 只改配置 config.yaml
  • 调用训练 → 只用 run.py
  • 部署模型 → 只看 deploy/ 目录

结构清晰、稳定、可扩展。

9.3 数据管线：并行加载 + 数据增强 + 高吞吐

CIFAR-10 数据管线模块负责：

  • 数据下载
  • 数据增强（随机裁剪、水平翻转）
  • 多进程加载
  • 训练/验证集区分

设计目标：高吞吐量 + 工程可维护。

9.4 模型模块化设计：SimpleCNN

模型定义模块 model.py 专注于一件事：

“给定输入 x，定义如何得到输出 logits。”

SimpleCNN 作为示例模型已足够清晰：Conv → ReLU → Conv → ReLU → Pool → FC → 分类。

模型作用：

  • 输入一张 32×32 RGB 图像
  • 输出 10 类的概率（飞机/汽车/猫/狗等 CIFAR10 类别）
  • 是一个可替换组件（可以改成 ResNet、MobileNet、ViT）
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class SimpleCNN(nn.Module):  
    """一个用于 CIFAR-10 的极简 CNN 网络。"""  
  
    def __init__(self, num_classes: int = 10):  
        super().__init__()  
  
        # 卷积特征提取模块：  
        # 输入：3x32x32  
        # 经过 Conv+ReLU+Conv+ReLU+MaxPool2d(2) 后，输出：64x16x16  
        self.conv = nn.Sequential(  
            nn.Conv2d(3, 32, kernel_size=3, padding=1),  # 输出 32x32x32  
            nn.ReLU(inplace=True),  
            nn.Conv2d(32, 64, kernel_size=3, padding=1), # 输出 64x32x32   
            nn.ReLU(inplace=True),  
            nn.MaxPool2d(2),                             # 下采样一半：64x16x16  
        )  
  
        # 全连接分类层：  
        # 将 64x16x16 展平成一个向量，然后映射到 num_classes        
        self.fc = nn.Linear(64 * 16 * 16, num_classes)  
  
    def forward(self, x):  
        # x: [B, 3, 32, 32]  
        x = self.conv(x)               # [B, 64, 16, 16]  
        x = x.view(x.size(0), -1)      # [B, 64*16*16]  
        x = self.fc(x)                 # [B, num_classes]  
        return x

9.5 Lightning 训练模块：训练逻辑从脚本中解放

在 lit_module.py 中，LightningModule 负责：

  • 前向传播
  • 训练步骤（loss 计算）
  • 验证步骤
  • 记录指标
  • 定义优化器

Lightning 自动处理：混合精度、多设备训练、日志回调、梯度累积、checkpoint 触发。工程训练流程大幅简化。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
class LitClassifier(L.LightningModule):  
    """将 nn.Module 包装成 LightningModule，托管训练细节。"""  
  
    def __init__(self, model, lr: float):  
        super().__init__()  
  
        # 真实的 PyTorch 模型（SimpleCNN）  
        self.model = model  
        # 学习率  
        self.lr = lr  
  
        # 将超参数保存到 checkpoint 中，方便恢复和可视化  
        self.save_hyperparameters(ignore=["model"])  
  
    def forward(self, x):  
        """推理 / 前向过程，直接调用内部模型。"""  
        return self.model(x)  
  
    def training_step(self, batch, batch_idx):  
        """单个训练 step 的逻辑。  
        Args:            
	        batch: 一个 batch 的数据，包含 (x, y)。  
            batch_idx: 当前 batch 的索引。  
        Returns:            
	        loss 张量，Lightning 会自动帮你做 backward。  
        """        
        x, y = batch  
        logits = self(x)  
        loss = F.cross_entropy(logits, y)  
  
        # self.log 会自动将指标记录到 logger（TensorBoard / W&B 等）  
        self.log("train_loss", loss, prog_bar=True, on_step=True, on_epoch=True)  
        return loss  
  
    def validation_step(self, batch, batch_idx):  
        """单个验证 step 的逻辑。"""  
        x, y = batch  
        logits = self(x)  
        val_loss = F.cross_entropy(logits, y)  
        self.log("val_loss", val_loss, prog_bar=True, on_step=False, on_epoch=True)  
  
    def configure_optimizers(self):  
        """配置优化器（和可选的学习率调度器）。"""  
        optimizer = optim.Adam(self.parameters(), lr=self.lr)  
        return optimizer

9.6 Hydra 配置：真正的可复现训练

整个训练系统的超参数都由配置管理：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
# 主配置文件：集中管理模型 / 数据 / 训练相关超参数  
# 可以通过命令行覆盖，例如：  
#   python run.py data.batch_size=128 train.lr=5e-4  
model:  
  name: simple_cnn  
  num_classes: 10          # CIFAR-10 有 10 个类别  
  
data:  
  batch_size: 64           # 每个 batch 的样本数  
  num_workers: 4           # DataLoader 后台进程数量  
  data_dir: ./data         # 数据集存放路径（会自动下载到这里）  
  
train:  
  max_epochs: 2           # 训练轮数  
  lr: 1e-3                 # 学习率  
  
  # 训练硬件 & 加速配置  
  precision: 32-true       # 默认用最稳定的 float32，可改为 16-mixed  
  accelerator: auto        # auto: 自动选择 gpu / mps / cpu  
  devices: 1               # 使用的设备个数（比如多卡时可以设置成 2 / 4）  
  
  # 梯度累积：显存较小时可以开启  
  accumulate_grad_batches: 1  
  
  # 是否对模型进行 torch.compile 加速（需要 PyTorch 2.0+）  
  compile: false  
  
  # 实验追踪（Weights & Biases），默认关闭  
  use_wandb: false  
  wandb_project: cifar10-project

我们可以轻松覆盖参数：

1
python run.py train.max_epochs=30 data.batch_size=128

不需要修改任何 Python 代码。

9.7 工程化训练：自动保存 + 自动导出模型

核心逻辑在 run.py：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
@hydra.main(config_path="configs", config_name="config", version_base=None)  
def main(cfg: DictConfig):  
    # 0. 设定随机种子，保证实验可复现  
    seed_everything(42)  
  
    # 1. 数据：只负责构建 DataLoader    
    train_loader, val_loader = build_dataloader(  
        batch_size=cfg.data.batch_size,  
        num_workers=cfg.data.num_workers,  
        data_dir=cfg.data.data_dir,  
    )  
  
    # 2. 模型：nn.Module + LightningModule  
    lit_model = build_model(cfg)  
  
    # 3. Trainer：Lightning 接管训练循环 / 日志 / AMP 等  
    trainer = build_trainer(cfg)  
  
    # 4. 启动训练  
    trainer.fit(lit_model, train_loader, val_loader)  
  
    # 5. 训练结束后：自动从 best.ckpt 导出 TorchScript 模型 model.pt    
    best_ckpt_path = os.path.join("checkpoints", "best.ckpt")  
    if os.path.exists(best_ckpt_path):  
        print(f"[INFO] 训练完成，发现 best.ckpt: {best_ckpt_path}")  
        print("[INFO] 正在从 best.ckpt 导出 TorchScript 模型为 model.pt ...")  
        export_torchscript(  
            checkpoint_path=best_ckpt_path,  
            out="model.pt",  
            num_classes=cfg.model.num_classes,  
        )  
        print("[INFO] 导出完成，可以直接使用 deploy/inference.py 或 deploy/api.py。")  
    else:  
        print("[WARN] 未找到 best.ckpt，跳过自动导出 model.pt。预期路径: checkpoints/best.ckpt")  

if __name__ == "__main__":  
    main()

训练日志（max_epochs先改为2）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
GPU available: True (mps), used: True
TPU available: False, using: 0 TPU cores
  | Name  | Type      | Params | Mode  | FLOPs
----------------------------------------------------
0 | model | SimpleCNN | 183 K  | train | 0    
----------------------------------------------------
183 K     Trainable params
0         Non-trainable params
183 K     Total params
0.733     Total estimated model params size (MB)
8         Modules in train mode
0         Modules in eval mode
0         Total Flops
Epoch 1: 100%|█████████████████████████████████████████████████████████████| 782/782 [01:03<00:00, 12.26it/s, v_num=1, train_loss_step=0.990, val_loss=1.070, train_loss_epoch=1.280]`Trainer.fit` stopped: `max_epochs=2` reached.                                               
Epoch 1: 100%|█████████████████████████████████████████████████████████████| 782/782 [01:03<00:00, 12.26it/s, v_num=1, train_loss_step=0.990, val_loss=1.070, train_loss_epoch=1.280]
[INFO] 训练完成，发现 best.ckpt: checkpoints/best.ckpt
[INFO] 正在从 best.ckpt 导出 TorchScript 模型为 model.pt ...
[OK] TorchScript 导出成功: model.pt
[INFO] 导出完成，可以直接使用 deploy/inference.py 或 deploy/api.py。

训练结束后自动生成：

  • best.ckpt（训练权重，Lightning 格式）
  • model.pt （TorchScript，可部署格式）

这一步完模型可以脱离 PyTorch 训练环境了，就是一个可以使用的模型了。

9.8 训练出的模型有什么用？

训练出的模型是“图像理解系统”的核心——一个 CIFAR-10 图像分类器。它能做什么？
JKr9SK

图 9-1 CIFAR-10数据集

输入一张如上图十分类的 RGB 图像，模型可自动识别其类别。后续可以换模型结构，换数据集，换任务（分类 → 检测/分割），达到不一样的目标。本质上这可作为视觉任务基底的工程框架。

9.9 模型部署：让模型成为真实服务

我们工程自带 deploy/ 目录，包含两个非常核心的部署模块。

9.9.1 本地推理（inference.py）

1
python deploy/inference.py --image xxx.jpg --model_path model.pt

例如：

1
2
3
4
5
6
// 输入一张青蛙图片
python deploy/inference.py --image test_image/gg.jpg --model_path model.pt  
[RESULT] 图片 test_image/gg.jpg 的预测类别 index 为: 6

python deploy/inference.py --image test_image/car.jpg --model_path model.pt
[RESULT] 图片 test_image/car.jpg 的预测类别 index 为: 1

功能：

  • 加载 TorchScript 模型
  • 自动预处理图像
  • 输出预测类别

用途：

  • 快速验证部署模型是否正常
  • 开发测试
  • 离线批量推理

9.9.2 在线 API 推理（FastAPI）

运行：

1
2
3
4
5
6
7
8
9
uvicorn deploy.api:app --host 0.0.0.0 --port=8000 --reload

INFO:     Will watch for changes in these directories: ['/xxx/dl_cifar10_project_autosave_export']
INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
INFO:     Started reloader process [40746] using StatReload
INFO:     Started server process [40748]
INFO:     Waiting for application startup.
[INFO] 模型已加载: model.pt
INFO:     Application startup complete.

启动一个 REST 推理服务。

1
2
curl -X POST -F "file=@test_image/gg.jpg" http://localhost:8000/predict
{"prediction":6}%

POST /predict

  • 上传一张图片
  • 返回分类结果
  • 前端 / APP / 其它服务可以直接调用
  • TorchScript 使推理速度快且稳定

这是项目部署的标准形式。

9.10 导出 ONNX：让模型跨平台

deploy/export.py 还能导出 ONNX：

1
2
3
4
python deploy/export.py

[OK] ONNX 导出成功: model.onnx
[OK] TorchScript 导出成功: model.pt

即可得到：

  • model.onnx
  • model.pt

ONNX 可以部署在：

  • TensorRT（NVIDIA）
  • ONNX Runtime
  • OpenVINO
  • iOS / Android / Web
  • C++ 程序

这让模型可具备跨平台能力。

9.11 全流程工程图

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
Hydra 配置（config.yaml）
            ↓
run.py （统一调度：数据 + 模型 + Trainer）
            ↓
DataLoader （并行加载 + 数据增强）
            ↓
模型构建（SimpleCNN）
            ↓
LightningModule（训练逻辑）
            ↓
Trainer（AMP + compile + DDP + Logging）
            ↓
Checkpoint（自动保存 best.ckpt）
            ↓
自动导出 TorchScript（model.pt）
            ↓
本地推理（inference.py）
            ↓
在线 API（FastAPI）
            ↓
ONNX 导出 → 全平台部署

这个流程已经完全符合现代深度学习工程体系。

小结
这一章完成了整篇内容的工程化收束，我们已经看到 PyTorch 如何从零散的代码生长为系统化的训练框架；如何从单机上的一次性试验扩展为可复现、可追踪的实验体系；如何让原本松散的 Python 脚本演化为结构清晰、可维护、可协作的项目工程；如何将朴素的单机训练进一步提升到带有 AMP、compile 等现代加速机制的高性能训练；以及如何让学术环境中的原型模型真正走向生产部署，成为可交付的工程成果。

总结一句：真正的深度学习，不只在训练模型，而在构建系统。

10. 总结：从“写代码”到“构建系统”

PyTorch 提供的是一整套从数值到结构，再到工程与系统的完整途径。前九章的内容看似分散：张量、模块、梯度、调试、性能、工程……。但它们从根本上围绕同一件事展开：让一个可微的计算世界，逐步演化成一个可组织、可优化、可调试、可扩展的系统。 以下是对整篇文章的总结：

  1. 梯度如何让系统能够“学习”
    第一章展示了一个最小训练流程。在那条简单的路径中，数据进入模型，进行预测，产生误差，梯度反向流动，参数随之变化。这是深度学习的原型，一个系统只要能够感受偏差并做出修正，就具备了学习的基本能力。
  2. 基本单位张量
    第二章讲了 PyTorch 的基础结构。Tensor 并不是普通数组，它同时携带数值、设备信息以及梯度追踪能力。当 requires_grad 被打开，它便成为计算图中的一个“节点”。这些节点在运算时互相连接，形成一张可微的图。整个系统的可学习性就在这张图中不断流动。
  3. Module可嵌套组合的结构
    第三章将注意力从“计算”带到了“组织”，Module 是构建神经网络的基础单位，它能够容纳参数、定义前向路径、组合子模块。一个复杂的网络，不过是 Module 的层层嵌套，结构的存在，让网络能够表达更多层级、概念以及空间关系。
  4. 数据如何进入模型
    第四章讲述了数据流动，Dataset 负责告诉系统每个样本的意义，DataLoader 则负责把数据源源不断地传递给模型。它不仅是数据传输工具，还包含缓存、随机化、多线程等机制。
  5. 动态计算图让模型有了可塑性
    第五章讲了 PyTorch 的核心，动态图意味着计算结构由数据即时决定，每一次前向传播都会产生一条全新的计算路径。grad_fn 记录了每个节点的来路，backward 则沿着这张图反向行走，完成误差的传递。detach 可以切断图，no_grad 可以冻结图。这一切构成了系统最核心的机制：能够对自身的“计算历史”进行可控的回溯。
  6. 调试
    第六章讲了如何排查错误。梯度可能爆炸，也可能消失；loss 可能发散，也可能完全不动；某些层可能停止更新，某些算子可能破坏计算图。这些现象不是随机，而是线索。通过检测梯度、分析激活、检查图结构、监控数值，我们能够一步步定位问题。调试能力让系统从“能运行”走向“能稳定地成长”。
  7. 性能优化
    第七章把我们带入现代深度学习的核心挑战。混合精度减少显存并提升吞吐；编译加速让 Python 级的执行变成图级优化；kernel 融合减少调度开销；数据流水线优化提高输入带宽；profiler 帮助找出瓶颈所在。这一章告诉我们，智能并不会自动变快，需要借助底层优化才能释放模型的真正潜力。
  8. 工程生态
    第八章让深度学习从“实验”走向“工程”。Hydra 让配置井然有序；Lightning 抽离了繁琐的训练样板；W&B 记录模型的一切变化；checkpoint 保存状态，使训练能够中断与恢复；部署工具让模型走向现实世界。工程化不是额外负担，而是智能系统能够被维护、扩展、协作和复现的基础。
  9. 端到端项目：计算、结构、数据、优化、工程
    第九章将所有内容汇聚成一条贯穿全流程的路径。Hydra 管理配置，run.py 作为调度中心；DataLoader 提供数据；model.py 定义结构；LightningModule 描述训练逻辑；Trainer 控制优化过程，同时连接 AMP、compile、DDP 等性能模块；W&B 管理实验；最终产出一个可复现、可部署的系统。

如果全篇内容提炼成一句核心思想：一个深度学习系统，是由可微分的计算、可组合的结构、可调试的反馈、可扩展的性能和可管理的工程共同构成的。

11. 备注