上一篇讨论了 Autograd：Tensor 运算如何形成动态计算图，backward() 如何沿图传播梯度，以及梯度状态和计算图生命周期之间有什么关系。

但一个真实的模型不会只是几个散落的 Tensor 运算。它通常包含：

• 多层嵌套的网络结构；
• 需要训练的参数；
• 不参与梯度更新但属于模型状态的 Buffer；
• 优化器和优化器状态；
• 训练、验证和推理三种不同阶段；
• 数据集、采样器和 DataLoader；
• CPU 到 GPU 的数据搬运；
• checkpoint、恢复训练和混合精度。

PyTorch 用 nn.Module 把这些对象组织起来，再通过 Optimizer、Dataset 和 DataLoader 将它们连接成训练系统：

1. Module 树
2. Parameter / Buffer / state_dict
3. forward
4. Loss
5. backward
6. Optimizer
7. Parameter 更新

本文要回答的不是“如何调用 nn.Linear”，而是：

PyTorch 如何把模型结构、参数状态、数据管线和训练循环组织成一个可以保存、迁移、复用和扩展的系统？¹

一、总览：从模型对象到训练系统

1. 本文的主线

本文沿着上面那条从 Module 树到参数更新的链路展开：先看一个模型对象里到底包含什么，nn.Module 靠什么机制发现子模块、Parameter 和 Buffer，state_dict 如何把模型状态变成可保存、可迁移的快照；再进入训练系统的其余部分——训练 / 评估 / 推理三种状态、Optimizer、Dataset / Sampler / DataLoader、CPU 到 GPU 的数据搬运；然后把这些部件拼成一个完整训练循环，并补上混合精度、Hook 和 checkpoint 三个工程上绕不开的话题。

%% 图：本文的主线：Module 树（第二至五章）→ 三种状态与 Optimizer（六、七）→ 数据管线（八、九）→ 拼成训练循环（十）→ 混合精度 / Hook / checkpoint（十一至十三）
flowchart TB
    subgraph M["模型对象（二 ～ 五章）"]
        direction LR
        A["nn.Module 树<br/>__init__ 定静态结构<br/>forward 定数据流"] --> B["注册机制<br/>__setattr__ 把 Parameter /<br/>子 Module / Buffer 登记进内部字典"] --> C["state_dict<br/>属性路径 → Tensor 的快照"]
    end
    subgraph T["训练部件（六 ～ 九章）"]
        direction LR
        D["train() / eval()<br/>no_grad / inference_mode"] --> E["Optimizer<br/>拿 parameters() 用 .grad 更新"]
        F["Dataset → Sampler → DataLoader<br/>→ pin_memory / non_blocking 搬到 GPU"]
    end
    M --> T
    T --> G["完整训练循环（十章）<br/>取 batch → forward → loss → backward → step"]
    G --> H["工程补丁（十一 ～ 十三章）<br/>autocast / GradScaler · Hook · checkpoint"]

2. 本文的章节安排

二、从模型对象到训练系统

1. 一个模型包含什么？

一个最小的训练模型可以写成：

import torch
from torch import nn


class MLP(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int) -> None:
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.activation = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        hidden = self.fc1(x)
        hidden = self.activation(hidden)
        return self.fc2(hidden)

可以看出，模型的定义其实是由两个部分组成：

1. __init__： 定义有哪些子 Module、Parameter 和 Buffer 形成模型的静态对象结构

2. forward： 定义数据如何经过这些 Module 形成一次执行中的计算关系

对应到上面的 MLP：

init 定义的 Module 树：

• MLP
  • fc1: Linear
    • weight
    • bias
  • activation: ReLU
  • fc2: Linear
    • weight
    • bias

forward 定义的数据流是 x → fc1 → activation → fc2 → output。

因此，这个对象从训练系统的角度同时包含：

• MLP
  • Module 层级结构
  • Linear 的 Parameter
  • ReLU 的行为
  • forward 定义的数据流
  • state_dict
  • train / eval 状态
  • device / dtype 迁移能力
  • hook 和序列化协议

nn.Module 的价值不只是提供一个可以继承的基类，而是定义了一套模型对象协议。

TIPS 对于上面这个简单的 MLP，其实我们可以直接使用 Sequential：

class SequentialMLP(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int) -> None:
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)

Sequential 内部已经固化了“按顺序调用每个子模块”的 forward() 逻辑，因此对于这种没有分支、旁路和多输入的模型，它非常合适。使用 Sequential 并不意味着没有计算关系，只是把这部分固定流程交给了容器类。两个版本在这个简单例子中的计算结果等价，区别在于 forward() 的控制权：

2. Module 树：__init__ 定义的静态结构

模型通常是一个树状结构。对于上面的 MLP，__init__ 执行完成后，Module 树是：

• MLP
  • fc1: Linear
    • weight
    • bias
  • activation: ReLU
  • fc2: Linear
    • weight
    • bias

可以查看这棵树：

model = MLP(128, 256, 10)
print(model)

也可以查看命名模块和参数：

for name, module in model.named_modules():
    print(name, type(module).__name__)

for name, parameter in model.named_parameters():
    print(name, parameter.shape)

这里得到的 fc1、activation、fc2 是静态对象结构的一部分。它们会影响：

• 参数发现；
• state_dict 的 key；
• .to(device) 的递归迁移；
• train() 和 eval() 的递归传播；
• hook 的注册；
• 分布式包装；
• checkpoint 加载。

但 Module 树只说明“有哪些模块和状态”，还没有说明一次输入会按照什么顺序经过它们。这个关系由 forward() 定义。

3. forward() 定义动态数据流

__init__() 和 forward() 描述的是两个不同层次：

• __init__()
  • 定义可复用的模块、参数和 Buffer
  • 通常在模型构造时执行一次
• forward()
  • 定义输入、模块和 Tensor 运算之间的数据流
  • 每次调用都可能形成一条新的执行路径

对于当前的 MLP，forward() 定义的是：

1. x
2. fc1
3. activation
4. fc2
5. output

在 Eager Mode 下，forward 中的每个 Tensor 运算会立即执行；如果开启梯度记录，Autograd 还会为本次调用记录一张动态计算图。这里要区分两张图：

• Module 树
  • → init 中注册的模块、参数和 Buffer
  • → 主要描述模型有哪些组件
• Tensor / Autograd 计算图
  • → forward 本次执行产生的运算关系
  • → 主要描述数据如何流过这些组件

forward() 不只是把层按顺序排列，也可以表达分支、跳跃连接和条件路径：

class ResidualBlock(nn.Module):
    def __init__(self, width: int) -> None:
        super().__init__()
        self.fc1 = nn.Linear(width, width)
        self.fc2 = nn.Linear(width, width)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x + residual

它的 Module 树只有两个 Linear，但 forward 的数据流包含一条旁路：

        ┌──────────── residual ────────────┐
        │                                  │
x ─────┴→ fc1 → ReLU → fc2 ────────────── + → output

因此，更准确的说法是：

__init__() 定义模型的静态组件结构，forward() 定义一次执行中的数据流；后者在开启 Autograd 时还会形成当前输入对应的动态计算图。

4. Module 是模型的生命周期容器

%% 图：nn.Module 是模型的生命周期容器：它持有的状态、它的行为、进出的通道
flowchart TB
    A[nn.Module]
    B[子 Module]
    C[Parameter]
    D[Buffer]
    E[Forward 行为]
    F[训练 / 评估状态]
    G[state_dict]
    H[设备与 dtype 迁移]

    subgraph STATE["它持有的状态"]
        direction TB
        B --> C --> D
    end
    subgraph BEH["它的行为"]
        direction TB
        E --> F
    end
    subgraph IO["进出的通道"]
        direction TB
        G --> H
    end
    A --> STATE
    A --> BEH
    A --> IO
    linkStyle 0,1,2,3 stroke-width:0px;

后续所有训练行为，都建立在 Module 能够找到并管理这些对象的前提上。

三、nn.Module 与模块注册

1. 为什么需要注册？

先澄清一件事：第二章的 MLP 里写的 self.fc1 = nn.Linear(...)、self.fc2 = nn.Linear(...) 已经是注册了——不是”没注册的写法”。nn.Module 重写了 Python 的属性赋值（下一节的 __setattr__），凡是把一个 Parameter、一个子 Module 赋给 self.xxx，它就悄悄把这个对象登记进自己的内部字典（_parameters / _modules）。你什么都不用多做，注册在赋值的那一刻就发生了。

那为什么还要有”注册”这个概念？因为训练系统里有一大堆事是替你自动做的，它们都靠这份登记簿：

登记簿上没有的对象，这五件事全部对它视而不见。下面就是唯一一种常见的”漏登记”：把子模块放进一个普通 Python list——list 既不是 Parameter 也不是 Module，__setattr__ 只会把它当普通属性存起来，里面的两个 Linear 于是从登记簿上消失：

class BadMLP(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layers = [                       # 普通 list：__setattr__ 看到的是 list，不是 Module
            nn.Linear(10, 20),
            nn.Linear(20, 1),
        ]

    def forward(self, x):
        for layer in self.layers:             # forward 里照样能用——问题不在这里
            x = layer(x)
        return x


class GoodMLP(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layers = nn.ModuleList([         # ModuleList 自己是一个 Module，里面的 Linear 会被登记
            nn.Linear(10, 20),
            nn.Linear(20, 1),
        ])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

两者的 forward 一模一样、都能算出结果，差别只在登记簿：

bad, good = BadMLP(), GoodMLP()
print(sum(p.numel() for p in bad.parameters()))    # 0     ← 优化器拿到 0 个参数
print(sum(p.numel() for p in good.parameters()))   # 241   ← 10×20+20 + 20×1+1
print(bad.state_dict().keys())                     # odict_keys([])：checkpoint 里什么都没有
bad.to("cuda"); print(bad.layers[0].weight.device) # cpu   ← 没搬动，前向时报 device mismatch

用 BadMLP 训练的症状很典型：代码不报错、loss 一直不降——因为 AdamW(bad.parameters()) 拿到的是空列表，optimizer.step() 每步什么都没更新。

这就是注册机制的意义：

只有被 Module 登记过的对象，才能自动参与参数遍历、状态保存、设备迁移和训练生命周期管理。 直接赋值 Parameter / Module 会自动登记；装在 list / dict / tuple 里的不会——要用 ModuleList / ModuleDict（第 3、4 节），它们自己是 Module，所以既能像容器一样用，又会把里面的东西登记上。

下面几节会用到 Parameter（可训练的张量）和 Buffer（随模型保存、但不训练的张量，比如 BatchNorm 的均值）两个词，第四章再给完整定义，这里知道它们是登记簿上除子模块之外的另两类条目即可。

2. __setattr__() 的作用

当执行：

self.layer = nn.Linear(10, 20)

Module 的属性设置逻辑会识别右侧对象是不是：

• Parameter；
• Module；
• 普通对象。

概念上可以简化为：

1. 赋值 self.name = value
2. value 是 Parameter？ → 注册到 _parameters
3. value 是 Module？ → 注册到 _modules
4. 否则 → 作为普通属性保存

把 Buffer 也算进来，__setattr__() 的完整决策路径如下——绿色分支的对象会进入 Module 的内部字典，从而被框架“看见”；灰色分支的对象只是 Python 对象属性：

%% 图：__setattr__ 的决策路径：Parameter、Module、已登记的 Buffer 进入内部字典，其余只是普通属性
flowchart TB
    S["self.name = value"]
    Q1{"value 是 nn.Parameter？"}
    Q2{"value 是 nn.Module？"}
    Q3{"name 已在 _buffers 中？"}
    P["_parameters#91;name#93; = value<br/>可训练参数"]
    M["_modules#91;name#93; = value<br/>子模块"]
    B["_buffers#91;name#93; = value<br/>模型状态 Buffer"]
    D["__dict__#91;name#93; = value<br/>普通属性"]
    V["框架可见：<br/>parameters() / state_dict() / .to() / train()"]
    I["框架不可见：<br/>不进 parameters()、state_dict()<br/>.to(device) 不会迁移"]

    S --> Q1
    Q1 -->|"是"| P
    Q1 -->|"否"| Q2
    Q2 -->|"是"| M
    Q2 -->|"否"| Q3
    Q3 -->|"是"| B
    Q3 -->|"否"| D
    P --> V
    M --> V
    B --> V
    D --> I

    classDef visible fill:#e8f5e9,stroke:#2e7d32;
    classDef hidden fill:#eeeeee,stroke:#757575;
    classDef q fill:#fff8e1,stroke:#f9a825;
    class P,M,B,V visible;
    class D,I hidden;
    class Q1,Q2,Q3 q;

图里三个菱形都答”否”时走最下面的灰色分支：值只是存进 Python 对象的 __dict__，成为一个普通属性——self.counter = 0、self.config = cfg、以及上一节那个 self.layers = [...] 都在这里。它们照常可以在 forward 里读写，只是框架的五件自动化的事都不认识它们。

注意 Buffer 有两条登记路径：一是“名字已经被 register_buffer() 登记过”——之后 self.scale = new_tensor 会更新 Buffer 而不是变成普通属性；二是 value 的类型是 torch.nn.Buffer（2.5 起，module.py 里 __setattr__ 有一个 isinstance(value, Buffer) 分支）——self.scale = nn.Buffer(torch.ones(10)) 与 register_buffer 等价，写法上和 Parameter 对称。普通 Tensor 直接赋值仍然只是属性。

这也是为什么下面几种对象的行为不同：

self.weight = nn.Parameter(torch.randn(10, 20))
self.layer = nn.Linear(10, 20)
self.counter = 0

分别属于：

• Parameter → 参数
• Module：→ 子模块
• int：→ 普通属性

3. ModuleList

三种容器解决的都是同一个问题：子模块的数量或名字在写代码时不固定（层数由配置决定、有几个输出头由任务决定），没法一个个写 self.layer1 = ...、self.layer2 = ...；又不能用普通 list / dict，否则第 1 节的漏登记就发生了。容器自己是一个 Module，赋给 self.layers 时被登记，它再把里面的每个子模块登记到自己名下（名字是 layers.0、layers.1……或 dict 的键）。三者的选择只看你需要哪种访问方式：

当需要保存一组按顺序执行的子模块时，可以使用：

class StackedMLP(nn.Module):
    def __init__(self, dims: list[int]) -> None:
        super().__init__()
        self.layers = nn.ModuleList(
            nn.Linear(left, right)
            for left, right in zip(dims, dims[1:])
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x)
            x = torch.relu(x)
        return x

ModuleList 提供列表式访问，同时保留 Module 的注册语义。

4. ModuleDict

当子模块需要使用有意义的名字时，可以使用：

class MultiHead(nn.Module):
    def __init__(self, input_dim: int, output_dim: int) -> None:
        super().__init__()
        self.heads = nn.ModuleDict({
            "classification": nn.Linear(input_dim, output_dim),
            "regression": nn.Linear(input_dim, 1),
        })

    def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        return {
            name: head(x)
            for name, head in self.heads.items()
        }

5. Sequential

Sequential 是一种更强约束的 Module 容器：

model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 1),
)

它适合“前一个模块的输出直接作为下一个模块的输入”的线性结构。

如果 forward 需要：

• 多个输入；
• 多个输出；
• 分支；
• 跳跃连接；
• 条件判断；

则通常应该使用自定义 Module，而不是强行塞进 Sequential。

四、Parameter、Buffer 与模型状态

1. Parameter 是什么？

Parameter 是 Tensor 的特殊封装，用来表示：这是 Module 的可训练参数。

from torch import nn

weight = nn.Parameter(torch.randn(4, 2))
print(weight.requires_grad)  # True

当它被赋值为 Module 的属性时，会自动注册：

class LinearLike(nn.Module):
    def __init__(self, input_dim: int, output_dim: int) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.randn(output_dim, input_dim))
        self.bias = nn.Parameter(torch.zeros(output_dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x @ self.weight.t() + self.bias

model = LinearLike(4, 2)
print(list(model.named_parameters()))

2. 普通 Tensor 不会自动成为参数

class NotAParameter(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.weight = torch.randn(4, 2, requires_grad=True)

这里的 weight 可以参与 Autograd，但它不会自动出现在：

list(model.parameters())
model.state_dict()

这再次说明：

• requires_grad：是否参与梯度计算
• Parameter 注册：是否被 Module 和 Optimizer 发现

3. Buffer 是什么？

有些 Tensor 属于模型状态，但不是需要 Optimizer 更新的参数。例如 BatchNorm 的运行统计量：

class RunningMean(nn.Module):
    def __init__(self, size: int) -> None:
        super().__init__()
        self.register_buffer("running_mean", torch.zeros(size))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x - self.running_mean

Buffer 的典型特征是：

• 不会出现在 parameters() 中；
• 会出现在 state_dict() 中；
• 会随 .to(device) 迁移；
• 可以是持久化的，也可以设置 persistent=False。

self.register_buffer(
    "temporary_mask",
    mask,
    persistent=False,
)

4. Parameter、Buffer 和普通属性

这张表是理解模型状态的关键。

5. 为什么普通 Tensor 不会自动迁移？

class BadModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.scale = torch.ones(10)

model = BadModule().cuda()
print(model.scale.device)  # 仍可能是 cpu

因为普通属性没有经过 Module 的注册机制。需要把它注册为 Buffer：

class GoodModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.register_buffer("scale", torch.ones(10))

五、state_dict：模型状态的结构化快照

1. state_dict() 保存什么？

model = LinearLike(4, 2)
state = model.state_dict()

print(state.keys())

通常会得到：

• weight
• bias

一个容易踩的坑：state_dict() 返回的 Tensor 是参数本体的引用（detach() 过，但共享 Storage），不是拷贝——state = model.state_dict() 之后再训练一步，state["weight"] 跟着变。所以“先记下最好的一版权重、训完再恢复”必须 {k: v.clone() for k, v in state.items()}（或 copy.deepcopy），异步保存 checkpoint 也要先拷贝一份再交给后台线程；torch.save 是同步序列化，直接传 state_dict() 没问题。

嵌套模块会使用点号组织 key：

• layers.0.weight
• layers.0.bias
• layers.2.weight
• layers.2.bias

state_dict 是一个从名字到 Tensor 的映射：

1. 模块路径 + 参数 / Buffer 名称
2. state_dict key
3. Tensor 状态

把 Module 树和 state_dict 并排放在一起看会更直观：左边是嵌套的对象树，右边是按”模块路径 + 属性名”拍平后的 key。Parameter 和持久化 Buffer 都会生成 key，无状态模块（如 ReLU）什么也不产生；下面以一个加了 BatchNorm 的 MLP 为例：

Module 树（嵌套对象）                      state_dict（扁平 key -> Tensor）
────────────────────────────────         ──────────────────────────────
MLP
├── fc1: Linear
│   ├── weight       (Parameter) ──────▶ "fc1.weight"
│   └── bias         (Parameter) ──────▶ "fc1.bias"
├── act: ReLU                (无状态，不产生任何 key)
├── bn: BatchNorm1d
│   ├── weight       (Parameter) ──────▶ "bn.weight"
│   ├── bias         (Parameter) ──────▶ "bn.bias"
│   ├── running_mean (Buffer)    ──────▶ "bn.running_mean"
│   ├── running_var  (Buffer)    ──────▶ "bn.running_var"
│   └── num_batches_tracked (Buffer) ─▶ "bn.num_batches_tracked"
└── fc2: Linear
    ├── weight       (Parameter) ──────▶ "fc2.weight"
    └── bias         (Parameter) ──────▶ "fc2.bias"

不包含：类定义 / forward 逻辑 / optimizer 状态 / scheduler / RNG 状态

2. state_dict 不是完整模型

state = model.state_dict()

它主要保存参数和持久化 Buffer，不包含：

• Python 类定义；
• forward 代码；
• Optimizer 状态；
• scheduler 状态；
• 随机数状态；
• 数据加载位置；
• 训练进度。

因此，模型状态快照 ≠ 可独立运行的完整模型。

3. 保存和加载模型状态

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

model = LinearLike(4, 2)
state = torch.load("model.pt", map_location="cpu")
model.load_state_dict(state)

加载时应该关注：

• key 是否一致；
• shape 是否一致；
• dtype 是否兼容；
• device 是否正确；
• 是否使用 strict 模式。

result = model.load_state_dict(state, strict=True)
print(result)

strict=False 可以允许缺失或多余 key，但不应该用来掩盖模型结构错误。

torch.load 底层是 pickle，反序列化可以执行任意代码。PyTorch 2.6 起 weights_only=True 成为默认值：只允许加载 Tensor、基本容器和显式加入白名单的类型。只保存 state_dict 的做法天然满足这个限制；如果 checkpoint 里放了自定义对象，在新版本上会遇到 UnpicklingError。这类版本行为的变化和 state_dict 自身的版本迁移，第十篇会作为兼容性问题专门讨论。

4. 保存可恢复训练的 checkpoint

如果需要中断后继续训练，通常需要保存更多状态：

checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "epoch": epoch,
    "global_step": global_step,
}

torch.save(checkpoint, "checkpoint.pt")

更完整的训练系统还可能保存：

• AMP GradScaler；
• Python 随机状态；
• NumPy 随机状态；
• PyTorch CPU 随机状态；
• CUDA 随机状态；
• 数据采样器状态；
• 配置和代码版本。

5. Java 对照的边界

state_dict 可以类比为结构化状态快照，但它不是 Java 序列化整个对象图：

• state_dict
  • → 只保存可训练参数和模型状态
  • → 依赖代码重新构造 Module
  • → 通过 key 和 shape 恢复

这种设计使得权重状态和模型代码相对解耦，也使得模型结构变更时必须显式处理兼容性。

六、训练、评估与推理状态

1. train() 与 eval()

model.train()
model.eval()

这两个方法会切换 Module 的训练状态，并递归作用于子模块。

典型受影响的模块包括：

• Dropout；
• BatchNorm；
• 某些自定义模块。

model.train()
output = model(inputs)

model.eval()
with torch.inference_mode():
    output = model(inputs)

2. eval() 不等于关闭梯度

model.eval()
output = model(inputs)
loss = criterion(output, targets)
loss.backward()

即使在 eval() 状态下，如果梯度记录开启，Autograd 仍然可以构建计算图。

反过来：

model.train()
with torch.no_grad():
    output = model(inputs)

训练状态下也可以临时关闭梯度记录。

两者控制的是不同维度：

• train / eval：Module 行为
• no_grad / inference_mode：Autograd 记录

两个维度正交，组合起来一共六种状态。下表每个格子写的是：Dropout / BatchNorm 的行为、是否构建计算图、典型用途：

纯推理的默认选择是右下角那一格：eval() 负责让 Module 行为确定，inference_mode() 负责让 Autograd 彻底退出。中间那格 eval() + no_grad() 同样是合法的推理写法，只是少省一点开销，而且产物之后还能参与求导——如果推理结果要喂给别的可微计算（蒸馏、评分器回传），就该用它。少了 eval() 才是错误：Dropout 会随机丢弃、BN 会用 batch 统计（MC-Dropout 之类刻意保留随机性的场景除外）。

3. 一个完整的评估函数

def evaluate(
    model: nn.Module,
    loader,
    criterion: nn.Module,
    device: torch.device,
) -> float:
    model.eval()
    total_loss = 0.0
    total_count = 0

    with torch.inference_mode():
        for inputs, targets in loader:
            inputs = inputs.to(device)
            targets = targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)

            batch_size = inputs.size(0)
            total_loss += loss.item() * batch_size
            total_count += batch_size

    return total_loss / total_count

这里同时处理了：

• Module 评估状态；
• 不记录梯度；
• 输入设备迁移；
• loss 转 Python 标量；
• 按样本数量汇总指标。

七、Optimizer 与参数更新

1. Optimizer 管理什么？

Optimizer 不只是一个“调用 step() 的对象”，它通常管理：

• 参数引用；
• 学习率；
• 参数组；
• 动量；
• 一阶和二阶统计量；
• weight decay 状态；
• step 计数。

model = LinearLike(4, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

训练过程为：

1. Parameter
2. forward
3. loss
4. backward
5. Parameter.grad
6. optimizer.step()
7. Parameter.data 更新

2. Optimizer 不负责计算梯度

loss.backward() 负责根据计算图计算梯度；optimizer.step() 负责根据已有梯度和自身状态更新参数。

二者职责不同：

• Autograd：→ 计算应该往哪个方向变化
• Optimizer → 决定采用什么更新规则和步长

3. 梯度清零

典型训练循环为：

optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()

如果不清零，梯度会按照上一篇介绍的规则继续累积。

• 第 1 步梯度 → grad
• 第 2 步梯度 → grad + 新梯度

4. 参数组

不同参数可以使用不同的学习率和 weight decay：

optimizer = torch.optim.AdamW([
    {
        "params": model.backbone.parameters(),
        "lr": 1e-5,
    },
    {
        "params": model.head.parameters(),
        "lr": 1e-3,
    },
])

参数组常用于：

• 微调预训练模型；
• 冻结部分层；
• 对 bias 和 normalization 参数使用不同正则化；
• 新增 head 使用更大学习率。

5. 冻结参数

for parameter in model.backbone.parameters():
    parameter.requires_grad_(False)

冻结参数意味着它们通常不参与梯度计算，但还需要注意：

• 是否仍然被放入 Optimizer；
• 模块是否处于 train 状态；
• BatchNorm 的 Buffer 是否仍然更新；
• forward 是否仍然消耗计算资源。

冻结参数不等于冻结整个模块的所有行为。

6. Optimizer state

AdamW 等 Optimizer 会为参数保存额外状态，可以通过 optimizer.state_dict() 查看。这些状态可能包括：

• step；
• exp_avg；
• exp_avg_sq。

因此训练显存通常不仅是“参数 + 梯度”，而是“参数 + 梯度 + Optimizer State”。

以 FP32 + Adam(W) 为例，每一个标量参数在显存里对应四块同样大小的内存，其中一半属于 Optimizer state：

这也是大模型训练中 Optimizer state 可能成为显存主要消耗者的原因之一：Adam 训练下每个参数要占 16 字节（参数、梯度、两个动量），第八篇会算这笔账。多卡训练时 Optimizer state 是最先被切分到各卡上的状态——第九篇的 ZeRO 与 FSDP 从这里开始。

八、Dataset、Sampler 与 DataLoader

1. Dataset 的职责

一个 Map-style Dataset 通常实现：

from torch.utils.data import Dataset


class RegressionDataset(Dataset):
    def __init__(self, features: torch.Tensor, targets: torch.Tensor) -> None:
        self.features = features
        self.targets = targets

    def __len__(self) -> int:
        return self.features.size(0)

    def __getitem__(self, index: int):
        return self.features[index], self.targets[index]

它提供两个基本协议：

• __len__：数据集大小
• __getitem__：根据索引取得一个样本

Dataset 不负责：

• 模型 forward；
• loss 计算；
• 参数更新；
• GPU 训练。

2. IterableDataset

对于流式数据或无法随机访问的数据，可以实现 IterableDataset：

from torch.utils.data import IterableDataset


class StreamingDataset(IterableDataset):
    def __iter__(self):
        for record in stream_records():
            yield encode(record)

它适合：

• 日志流；
• 大于本地磁盘的数据；
• 在线生成的数据；
• 顺序读取的数据源。

但多 worker 场景需要显式处理 worker 分片，否则多个 worker 可能重复读取相同数据。

3. Sampler 的职责

Sampler 决定索引访问顺序：

• Dataset：提供样本
• Sampler：产生索引顺序
• DataLoader：组织 batch 和 worker

常见 Sampler 包括：

• SequentialSampler；
• RandomSampler；
• WeightedRandomSampler；
• DistributedSampler。

数据顺序会影响：

• 训练随机性；
• 类别平衡；
• 分布式数据切分；
• 可复现性。

4. DataLoader 的职责

from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
)

DataLoader 通常负责：

• 读取样本；
• 组织 batch；
• 调用 collate_fn；
• 启动 worker；
• 预取数据；
• 将 CPU 数据放入 pinned memory。

5. collate_fn

默认 collate 可以处理形状一致的样本：

样本 1：shape=(3,)
样本 2：shape=(3,)
样本 3：shape=(3,)
    ↓
Batch：shape=(3, 3)

但变长序列通常需要自定义 padding：

def collate_fn(batch):
    sequences, labels = zip(*batch)
    padded = pad_sequence(
        sequences,
        batch_first=True,
        padding_value=0,
    )
    return padded, torch.tensor(labels)

6. DataLoader 为什么会成为瓶颈？

训练的一次迭代可以表示为：

1. 磁盘读取
2. 数据解码
3. CPU 预处理
4. collate
5. Pinned Memory
6. CPU → GPU
7. GPU 计算

把这条链路展开到进程 / 线程层面，就能看到 num_workers、prefetch_factor、pin_memory 三个参数各自作用在哪一段，以及瓶颈最终在哪里表现出来：

%% 图：DataLoader 的进程 / 线程链路：num_workers、prefetch_factor、pin_memory 各作用在哪一段
flowchart TB
    subgraph main["DataLoader 主进程"]
        SM["Sampler<br/>产生本 batch 的索引列表"]
        IQ["index_queue<br/>把索引分发给 worker"]
        RQ["结果队列（worker_result_queue）<br/>深度 = num_workers × prefetch_factor"]
        PIN["pin_memory 线程<br/>把 batch 复制到 pinned memory"]
        IT["for inputs, targets in loader:<br/>训练循环取走 batch"]
    end
    subgraph workers["worker 进程池（num_workers 个）"]
        GI["Dataset.__getitem__(idx)<br/>读盘、解码、预处理（CPU 密集）"]
        CO["collate_fn<br/>把样本堆叠成 batch Tensor"]
    end
    H2D["inputs.to(device, non_blocking=True)<br/>CPU -> GPU 复制（PCIe）"]
    GPU["GPU forward / backward / step"]
    WAIT["瓶颈点：GPU 等 batch<br/>结果队列为空 -> GPU 空转"]

    SM --> IQ
    IQ --> GI
    GI --> CO
    CO --> RQ
    RQ --> PIN
    PIN --> IT
    IT --> H2D
    H2D --> GPU
    RQ -.->|"队列空"| WAIT
    WAIT -.-> GPU

    classDef knob fill:#fff8e1,stroke:#f9a825;
    classDef par fill:#e3f2fd,stroke:#1565c0;
    classDef gpu fill:#e8f5e9,stroke:#2e7d32;
    classDef bad fill:#ffebee,stroke:#c62828;
    class RQ,PIN knob;
    class GI,CO par;
    class GPU,H2D gpu;
    class WAIT bad;

• num_workers 决定蓝色部分有多少个进程并行执行 __getitem__ + collate_fn；
• prefetch_factor（默认 2）决定每个 worker 最多提前准备多少个 batch，也就是结果队列的深度；
• pin_memory=True 才会启动 pin_memory 线程，后面的 non_blocking=True 复制才可能真正异步。

如果前面的数据准备速度低于 GPU 消耗速度，就会出现：

1. GPU 计算完成
2. 等待下一个 batch
3. GPU 利用率下降

所以 GPU 利用率低不一定是模型 Kernel 的问题，也可能是 DataLoader 没有及时提供数据。第八篇把这种情况单列为一类瓶颈（GPU 在等数据到达），并给出在 Profiler 时间线上把它与”GPU 在等 Python”区分开的方法。

7. num_workers 不是越大越好

增加 worker 可能提升并行数据准备能力，但也会增加：

• 进程数量；
• CPU 内存；
• 进程间通信；
• 文件句柄；
• 数据复制或序列化；
• worker 启动时间。

正确的调优方式是测量：

num_workers=0
    ↓
num_workers=2
    ↓
num_workers=4
    ↓
num_workers=8

观察：

• batch 等待时间；
• GPU 利用率；
• CPU 使用率；
• 主机内存；
• 端到端 step time。

8. DataLoader 的 Java 对照

可以把数据管线近似理解为生产者—消费者系统：

• Dataset / Worker：生产样本
• Prefetch Queue：暂存 batch
• Training Loop：消费 batch
• GPU：消费计算输入

但 DataLoader 不是简单的 Java ExecutorService：它还涉及 Python 进程、Tensor 共享、pinned memory 和设备搬运。

九、CPU-GPU 数据传输与训练流水线

1. 模型和数据必须位于兼容设备

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

for inputs, targets in loader:
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)

.to(device) 需要分别作用于：

• 模型参数；
• 模型 Buffer；
• 输入 Tensor；
• target Tensor；
• 运行时创建的辅助 Tensor。

2. 模型迁移的递归性

model.cuda() 会递归迁移已注册的 Parameter 和 Buffer，以及子 Module 中的对应状态。

但不会自动迁移普通属性中的 Tensor：

self.mask = torch.ones(10)

这正是 Buffer 注册重要的原因。

3. Pinned Memory

CPU 普通内存与 GPU 之间的数据复制，可能受到主机内存页锁定状态的影响。DataLoader 可以配置：

loader = DataLoader(
    dataset,
    batch_size=64,
    pin_memory=True,
)

再配合：

inputs = inputs.to("cuda", non_blocking=True)

在满足条件时，数据复制可以与其他 GPU 工作重叠。

但需要注意：

• pin_memory=True 会增加主机内存压力；
• non_blocking=True 不保证所有场景都异步；
• 如果后面立刻读取结果，仍可能发生同步；
• 数据预处理和 PCIe 带宽仍可能是瓶颈。

4. 训练流水线的理想状态

CPU Worker 准备 Batch N+1
              │
              ▼
GPU 计算 Batch N
              │
              ▼
CPU Worker 准备 Batch N+2

目标是让 CPU 数据准备和 GPU 计算尽量重叠，而不是：

1. CPU 准备 Batch N
2. GPU 计算 Batch N
3. CPU 准备 Batch N+1
4. GPU 等待

第八篇性能文章会用 Profiler 分析这种等待具体发生在哪里。

5. 不要在训练循环中频繁搬回 CPU

for batch in loader:
    output = model(batch)
    value = output.cpu().numpy()

如果每个 step 都把大 Tensor 搬回 CPU，可能导致：

• GPU 与 CPU 同步；
• PCIe 带宽消耗；
• Python 线程等待；
• 显存和主机内存之间频繁交换。

如果只是记录标量，应优先使用：

value = loss.detach().item()

十、完整训练循环

1. 最小训练循环

def train_one_epoch(
    model: nn.Module,
    loader,
    criterion: nn.Module,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
) -> float:
    model.train()
    total_loss = 0.0
    total_count = 0

    for inputs, targets in loader:
        inputs = inputs.to(device)
        targets = targets.to(device)

        optimizer.zero_grad(set_to_none=True)
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        batch_size = inputs.size(0)
        total_loss += loss.detach().item() * batch_size
        total_count += batch_size

    return total_loss / total_count

2. 一个完整的训练程序

model = MLP(input_dim=128, hidden_dim=256, output_dim=10)
criterion = nn.CrossEntropyLoss()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)                                              # 先迁移
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)    # 再构造 optimizer（官方文档对 Module.to/cuda 的提醒）

for epoch in range(10):
    train_loss = train_one_epoch(
        model,
        train_loader,
        criterion,
        optimizer,
        device,
    )
    valid_loss = evaluate(
        model,
        valid_loader,
        criterion,
        device,
    )
    print(epoch, train_loss, valid_loss)

3. 一次训练迭代的状态流

%% 图：一次训练迭代的状态流：DataLoader → Module → Autograd → Optimizer
sequenceDiagram
    participant L as DataLoader
    participant M as Module
    participant A as Autograd
    participant O as Optimizer

    L->>M: inputs, targets
    M->>M: forward(inputs)
    M->>A: 构建动态计算图
    A-->>M: outputs
    M->>A: loss.backward()
    A-->>O: parameter.grad
    O->>O: optimizer.step()
    O-->>M: 更新 Parameter
    M->>O: zero_grad()

4. 训练循环中的常见顺序错误

忘记调用 optimizer.step()

loss.backward()
# 没有 optimizer.step()

这样只计算了梯度，没有更新参数。

忘记清空梯度

loss.backward()
optimizer.step()

如果没有在下一步前清零，梯度会继续累积。

清零放在哪一行

通常推荐：

optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()

关键不是某一行绝对不能换位置，而是每一次 backward 前必须明确当前 .grad 中应该保留什么。

验证阶段仍然构建计算图

model.eval()
for batch in valid_loader:
    output = model(batch)

这可能产生不必要的 Autograd 和显存开销。验证通常应使用：

model.eval()
with torch.inference_mode():
    ...

十一、混合精度训练入门

1. 为什么使用混合精度？

低精度计算可能带来：

• 更小的存储空间；
• 更低的内存带宽需求；
• 更高的硬件吞吐；
• 更低的通信成本。

但低精度也可能带来：

• 溢出；
• 下溢；
• 梯度数值不稳定；
• 不同算子支持差异。

2. autocast

现代 PyTorch 中可以使用：

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    outputs = model(inputs)
    loss = criterion(outputs, targets)

autocast 会根据算子和设备选择适合的计算 dtype。它不是把整个模型的所有 Tensor 无条件转换成一种类型。

3. GradScaler

对于 FP16 训练，通常还需要梯度缩放：

scaler = torch.amp.GradScaler("cuda")

for inputs, targets in loader:
    optimizer.zero_grad(set_to_none=True)

    with torch.autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = criterion(outputs, targets)

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

这三行各自做了什么、和 autocast、Optimizer 之间怎么配合，用一次迭代的时序来看：

%% 图：GradScaler 的一次迭代：scale 放大 loss，step 内部 unscale 并查溢出，update 自适应 scale
sequenceDiagram
    participant T as 训练循环
    participant AC as autocast
    participant S as GradScaler
    participant O as Optimizer

    T->>AC: with autocast(dtype=float16)
    activate AC
    Note over AC: forward 中 matmul 等算子用 FP16 计算<br/>loss 等归约算子保持 FP32
    AC-->>T: outputs, loss
    deactivate AC

    T->>S: scaler.scale(loss)
    S-->>T: loss × scale（scale 初始 65536）
    T->>T: (loss × scale).backward()
    Note over T: 梯度被同比例放大<br/>避免 FP16 下小梯度下溢成 0

    T->>S: scaler.step(optimizer)
    activate S
    S->>S: unscale_(optimizer)：grad ÷= scale
    S->>S: 检查所有 grad 是否含 inf / nan
    alt 梯度正常
        S->>O: optimizer.step()
        O-->>S: 参数已更新
    else 发现 inf / nan
        Note over S,O: 跳过本步 optimizer.step()<br/>参数保持不变
    end
    deactivate S

    T->>S: scaler.update()
    Note over S: 本步溢出 -> scale ÷= 2<br/>连续 2000 步正常 -> scale ×= 2

要点有两个：scaler.step() 内部才做 unscale 和溢出检查，所以如果需要在 step 前做梯度裁剪，必须先显式调用 scaler.unscale_(optimizer)；scaler.update() 让 scale 自适应——溢出就减半、长期稳定就翻倍，这也是它必须在每步 step() 之后调用的原因。

BF16 的指数范围接近 FP32，很多训练场景不需要和 FP16 完全相同的缩放策略，但具体行为仍取决于硬件、算子和版本。

4. 混合精度的边界

混合精度不是“免费加速”：

• 首次运行可能有额外上下文开销；
• 某些算子可能回退到 FP32；
• dtype 转换本身需要成本；
• 数值结果可能发生变化；
• 不同 GPU 的收益不同。

应使用正确性测试和 Benchmark 验证，而不是只看代码中出现了 autocast。本节只介绍混合精度作为训练 API 的用法；它为什么能加速（Tensor Core 与访存量减半）、什么时候加速不明显、如何用 Benchmark 确认收益，第八篇作为一条性能处方展开。

十二、Hooks 与模型观测

1. Hook 能做什么？

Module 提供多种 hook，用于观察或干预执行过程：

def log_shape(module, inputs, output):
    print(type(module).__name__, output.shape)

handle = model.register_forward_hook(log_shape)

使用结束后应移除：

handle.remove()

Hook 常用于：

• 观察中间激活；
• 调试 shape；
• 统计耗时；
• 采集特征；
• 检查模块执行顺序；
• 研究模型行为。

2. Hook 的成本和风险

Hook 会改变或包裹执行路径，可能带来：

• Python 回调开销；
• 额外 Tensor 引用；
• 计算图被意外保留；
• 编译器 graph break；
• 并发和生命周期复杂化。

因此，不应在生产热路径中无条件挂载大量 hook。

3. Hook 与 forward() 的关系

output = model(inputs) 通常会经过 Module 的调用协议，而不是简单执行 model.forward(inputs)。Hook 正是 Module 调用协议中的一部分。直接调用 forward() 可能绕过部分 Module 行为，因此业务代码通常应调用 model(inputs)，而不是 model.forward(inputs)。这与第一篇对 PyTorch 调用路径的介绍相互呼应。

十三、Checkpoint 与可恢复训练

1. 只保存模型权重

适合部署或只需要推理的场景：

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

2. 保存完整训练状态

适合训练中断后恢复：

checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "epoch": epoch,
    "global_step": global_step,
}

如果使用 scheduler 和 AMP，还应加入：

checkpoint["scheduler"] = scheduler.state_dict()
checkpoint["scaler"] = scaler.state_dict()

3. 保存 checkpoint 的时机

常见选择包括：

• 每个 epoch 结束；
• 每隔固定 global step；
• 验证集指标刷新时；
• 训练发生异常前；
• 周期性保存最新和最佳版本。

生产系统还需要考虑：

• 原子写入；
• 临时文件；
• 文件损坏；
• 并发保存；
• 远程对象存储；
• 版本和元数据；
• 清理旧 checkpoint。

4. Resume 不只是加载权重

checkpoint = torch.load(path, map_location="cpu")
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
start_epoch = checkpoint["epoch"] + 1

如果只加载模型而不加载 Optimizer state，训练可能不能从原来的优化轨迹继续。

一个完整的 save / resume 流程如下。save 时是“收集所有会影响后续训练轨迹的状态”，resume 时则有严格的先后顺序——模型对象必须先由代码构造出来，optimizer 又必须在 model.to(device) 之后才能 load_state_dict，否则 Optimizer state 会留在 CPU 上、与参数设备不一致：

%% 图：save 时收集的全部状态：模型、优化器、scheduler、scaler、进度、RNG 与 Sampler 位置
flowchart TB
    subgraph trainer["训练器持有的状态（save 时收集）"]
        S1["model.state_dict()<br/>参数 + 持久化 Buffer"]
        S2["optimizer.state_dict()<br/>exp_avg / exp_avg_sq<br/>step / param_groups"]
        S3["scheduler.state_dict()<br/>last_epoch、base_lrs"]
        S4["scaler.state_dict()<br/>当前 scale、growth_tracker"]
        S5["epoch + global_step<br/>训练进度"]
        S6["RNG 状态<br/>random / numpy / torch / cuda"]
        S7["DataLoader / Sampler 位置<br/>epoch 内已消费的 batch 数"]
        S1 ~~~ S2
        S3 ~~~ S4
        S5 ~~~ S6
    end
    CK["checkpoint = dict(...) 收齐全部状态<br/>torch.save(checkpoint, tmp)<br/>os.replace(tmp, path) 原子写入"]
    S2 --> CK
    S4 --> CK
    S6 --> CK
    S7 --> CK

    classDef st fill:#e3f2fd,stroke:#1565c0;
    classDef io fill:#fff8e1,stroke:#f9a825;
    class S1,S2,S3,S4,S5,S6,S7 st;
    class CK io;

resume 是 save 的逆过程，但顺序不能乱：

%% 图：resume 的六步顺序：先构造对象，load 权重，to(device)，再 load 优化器状态
flowchart TB
    R0["torch.load(path, map_location='cpu')<br/>weights_only=True 默认，只含 Tensor / 基本容器"]
    R1["1. 用代码重新构造 model / optimizer / scheduler<br/>state_dict 里没有类定义和 forward"]
    R2["2. model.load_state_dict(ckpt#91;'model'#93;, strict=True)<br/>校验 key 与 shape"]
    R3["3. model.to(device)"]
    R4["4. optimizer.load_state_dict(ckpt#91;'optimizer'#93;)<br/>必须在 to(device) 之后：state 会被搬到参数所在设备"]
    R5["5. scheduler / scaler.load_state_dict"]
    R6["6. 恢复 RNG 状态，start_epoch = epoch + 1<br/>Sampler.set_epoch / 跳过已消费 batch"]
    R0 --> R1 --> R2 --> R3 --> R4 --> R5 --> R6

    classDef io fill:#fff8e1,stroke:#f9a825;
    classDef step fill:#e8f5e9,stroke:#2e7d32;
    classDef warn fill:#ffebee,stroke:#c62828;
    class R0 io;
    class R1,R2,R3,R5,R6 step;
    class R4 warn;

这就像 Java 服务恢复业务状态时，不只是恢复一个对象字段，还可能需要恢复：

• 定时任务进度；
• 消费位点；
• 缓存状态；
• 重试计数；
• 外部依赖版本。

十四、Java 工程师如何理解 nn.Module

1. Module 更像带状态协议的组件树

nn.Module 可以近似理解为：

组件树
+ 参数注册
+ 状态注册
+ 生命周期切换
+ 设备迁移
+ 序列化协议
+ 调用钩子

它不是简单的 Java POJO，也不是只有一个 forward() 方法的接口。

2. Parameter 不是普通字段

Java 中一个 double[] 字段不会因为被放入对象就自动进入 Optimizer。PyTorch 中，Parameter 的注册让框架能够发现它、保存它、迁移它并更新它。

• 普通 Tensor 字段：只是对象属性
• nn.Parameter 字段：Module 管理的模型参数

3. state_dict 不是 Java Serialization

state_dict 保存的是显式模型状态，而不是整个 Python 对象图。重新加载时仍然需要：

1. 重新构造模型代码
2. 加载 state_dict
3. 恢复参数和 Buffer

这种方式牺牲了部分“直接恢复对象”的便利性，换来了更明确的状态边界和更好的跨代码版本控制能力。

4. DataLoader 是带设备语义的数据管线

DataLoader 可以类比生产者—消费者，但它还额外包含：

• Python worker 进程；
• Tensor batch；
• pinned memory；
• CPU-GPU 传输；
• 数据预取；
• 训练 step 的背压。

所以调优 DataLoader 不能只看线程池大小，还要看数据格式、CPU、内存、总线和 GPU 消费速度。

十五、本文小结

本文讨论了 PyTorch 如何把 Module、参数、状态、数据和训练循环组织成完整系统。

1. Module 是模型生命周期容器

• Module
  • 子 Module
  • Parameter
  • Buffer
  • forward
  • train / eval
  • state_dict
  • device / dtype 迁移
  • hooks

2. 注册机制决定框架能否发现对象

• 普通 Python list：不自动注册子 Module
• ModuleList / ModuleDict：注册子 Module
• 普通 Tensor 属性：不自动进入 state_dict，也不自动迁移
• register_buffer()：注册模型状态，但不参与 Optimizer 更新

3. 训练系统的状态流

1. DataLoader
2. inputs / targets
3. model(inputs)
4. loss
5. loss.backward()
6. Parameter.grad
7. optimizer.step()
8. Parameter 更新

4. 训练和推理的边界

• train()：训练态 Module 行为
• eval()：评估态 Module 行为
• no_grad()：关闭梯度记录
• inference_mode()：面向纯推理的更强上下文

5. 分析一个训练问题的顺序

1. Module 是否正确注册？
2. Parameter 是否被发现？
3. Buffer 是否正确迁移？
4. 输入和模型是否在同一 device？
5. DataLoader 是否及时供给 batch？
6. loss 是否连接到目标参数？
7. grad 是否正确清零和更新？
8. checkpoint 是否保存了完整状态？

6. 本篇涉及的源码位置

本篇讨论的机制在源码中的位置（对应第一篇第七章的代码地图）：

下一篇将进入 PyTorch 的算子运行时：

同一个 add 算子为什么能够运行在 CPU、CUDA、Autograd 和 Meta 等不同后端上？Dispatcher 又是如何选择具体实现的？

十六、自测

1. self.layers = [nn.Linear(4, 4) for _ in range(3)] 会有什么问题？parameters() 能看到它们吗？

   答案

   看不到——list 不是 Module，__setattr__ 不登记，参数不进优化器、不进 state_dict、.to("cuda") 也不迁移；要用 nn.ModuleList（或 nn.Sequential）。

2. register_buffer("running_mean", t) 与 self.running_mean = t（普通 Tensor 属性）差在哪？

   答案

   buffer 进 state_dict（能保存 / 加载）、随 .to() 迁移、随 train()/eval() 无关但可被 hook 访问；普通属性都不会，换设备后停在 CPU 上、checkpoint 里也没有。persistent=False 的 buffer 迁移但不保存。

3. load_state_dict(sd, strict=False) 会隐藏什么错误？什么时候该用？

   答案

   缺失的键（模型里有、文件里没有的参数保持随机初始化）与多余的键都只记录不报错——最常见的“加载了 checkpoint 但效果像没训过”的来源；只在刻意加载部分权重（换了分类头）时用，并检查返回的 missing_keys / unexpected_keys。

4. torch.load(path) 默认为什么不安全？weights_only=True 做了什么？

   答案

   torch.save 用 pickle，任意 Python 对象（含 __reduce__ 里的代码）都能被反序列化执行——加载不可信文件等于执行任意代码；weights_only=True 用受限 unpickler 只允许 Tensor、基本容器等白名单类型。

5. DataLoader(num_workers=4) 的 4 个 worker 是线程还是进程？pin_memory=True 在哪个线程做、为什么能加速？

   答案

   进程（绕开 GIL 做 Python 预处理，每个 epoch 重新启动除非 persistent_workers）；pin 在主进程的一个专用线程里把 batch 拷到页锁定内存，之后 to("cuda", non_blocking=True) 才能真正异步 DMA，与计算重叠。

• 上一篇
  PyTorch 深度实践（03）：自动求导与动态计算图
• 下一篇
  PyTorch 深度实践（05）：Dispatcher 与算子系统