第9章 nn.Module：参数注册、Hooks、序列化

"Calls super().__setattr__('a', a) instead of the typical self.a = a to avoid Module.__setattr__ overhead. Module's __setattr__ has special handling for parameters, submodules, and buffers but simply calls into super().__setattr__ for all other attributes."

—— torch/nn/modules/module.py，Module.__init__ 内的注释（v2.11.0）

9.1 一段神奇的代码

每个 PyTorch 用户都写过这样的模型：

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(10, 20)
        self.linear2 = nn.Linear(20, 30)
        self.bn = nn.BatchNorm1d(20)
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        x = self.linear1(x)
        x = self.bn(x)
        x = self.linear2(x)
        return x

model = MyModel()
print(list(model.parameters()))   # 自动列出 linear1.weight, linear1.bias, linear2.weight, ..., bn.weight, bn.bias
print(model.state_dict().keys())   # 自动收集 'linear1.weight', 'bn.running_mean', ...
torch.save(model.state_dict(), 'ckpt.pt')   # 序列化

魔法到处都是：

1. self.linear1 = nn.Linear(...) 怎么自动加入 model.parameters()？
2. bn.running_mean（不是 Parameter）怎么也进了 state_dict？
3. state_dict() 的 key 'linear1.weight' 是怎么拼出来的？
4. model(x) 调 forward(x) 还是别的？

这些问题的答案都在 torch/nn/modules/module.py（3054 行）。本章拆它的设计。

9.1.1 nn.Module 的设计目标

Module 不只是"参数容器" —— 它是 PyTorch 模型的统一抽象，要解决几个问题：

• 递归组合：模型由层组成，层又由更小层组成。需要无限递归
• 参数自动管理：用户写 self.linear = nn.Linear(...) 自动让 optimizer 看到所有权重
• 状态保存与加载：训练 ckpt、模型部署都需要"把所有可变状态存到磁盘"
• 跨 device 迁移：.cuda() / .cpu() / .half() 要递归整个网络
• 训练 / 推理模式切换：BN / Dropout 等需要在 train / eval 时不同行为
• 可扩展性：第三方框架（HuggingFace、PyTorch Lightning、DDP 等）能 hook 进 Module 生命周期

这些目标共同塑造了 Module 的内部结构。理解每个设计选择背后的目标，能让你看 nn.Module 源码不再迷失在 3054 行细节里。

9.2 __init__：建立四个核心容器

打开 torch/nn/modules/module.py:483：

class Module:
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        ...
        super().__setattr__("training", True)
        super().__setattr__("_parameters", {})
        super().__setattr__("_buffers", {})
        super().__setattr__("_non_persistent_buffers_set", set())
        super().__setattr__("_backward_pre_hooks", OrderedDict())
        super().__setattr__("_backward_hooks", OrderedDict())
        super().__setattr__("_is_full_backward_hook", None)
        super().__setattr__("_forward_hooks", OrderedDict())
        super().__setattr__("_forward_hooks_with_kwargs", OrderedDict())
        super().__setattr__("_forward_hooks_always_called", OrderedDict())
        super().__setattr__("_forward_pre_hooks", OrderedDict())
        super().__setattr__("_forward_pre_hooks_with_kwargs", OrderedDict())
        super().__setattr__("_state_dict_hooks", OrderedDict())
        super().__setattr__("_state_dict_pre_hooks", OrderedDict())
        super().__setattr__("_load_state_dict_pre_hooks", OrderedDict())
        super().__setattr__("_load_state_dict_post_hooks", OrderedDict())
        super().__setattr__("_modules", {})

注意每行都用 super().__setattr__ 而不是 self.xxx = ... —— 这是为了绕过自己重写的 __setattr__（直接给 __dict__ 赋值）。开头引的那段源码注释说得很直白：为的是省掉 Module.__setattr__ 的开销；何况此刻 _parameters 等字典还不存在，Module.__setattr__ 里那套分派本就无从谈起。

四个核心容器：

字段	类型	用途
_parameters	dict[str, Parameter]	可学习参数（weight、bias 等）
_buffers	dict[str, Tensor]	不可学习但需保存的状态（如 BN 的 running_mean）
_non_persistent_buffers_set	set[str]	不存进 state_dict 的 buffer 名（register_buffer(..., persistent=False) 注册的）
_modules	dict[str, Module]	子模块（自动递归整个网络）

剩下的 12 个字段里有 11 个是 hook 容器（OrderedDict 保证调用顺序），另一个 _is_full_backward_hook 不是容器而是一个三态标志（bool | None，见 torch/nn/modules/module.py:462 的类型标注），记录这个 Module 上注册的是新版还是旧版 backward hook（torch/nn/modules/module.py:1449 用它分流）。再加上第一行的 training 标志（训练 / 推理模式切换改的就是它，见 §9.9），Module 的所有"高级特性"都建立在这 17 个字段之上。

为什么用 OrderedDict 而不是 list？因为 hook 注册返回 RemovableHandle，handle 持有一个 id，用户调 handle.remove() 时通过 id 从 dict 删除。dict 的随机访问 O(1) 比 list 的线性查找方便。OrderedDict 保留插入顺序让多个 hook 按注册顺序调用 —— 这是用户期望的语义。

至于 _parameters / _buffers / _modules 用普通 dict 而不是 OrderedDict：Python 3.7+ 普通 dict 已经保证插入顺序，不需要显式 OrderedDict。hook 容器仍用 OrderedDict 有个实打实的用处：register_forward_hook(..., prepend=True) 这类注册靠 OrderedDict.move_to_end(handle.id, last=False) 把新 hook 挪到最前（torch/nn/modules/module.py:1752），普通 dict 没有这个方法。

graph TB
    M[nn.Module 实例]
    M --> P["_parameters dict<br/>weight / bias / ..."]
    M --> B["_buffers dict<br/>running_mean / num_batches_tracked"]
    M --> Mods["_modules dict<br/>子模块递归"]
    M --> H["11 个 hook 相关 OrderedDict<br/>前/后/状态/反向"]
    Mods --> Sub1[子 Module 1]
    Mods --> Sub2[子 Module 2]
    Sub1 --> SubP1[子 Module 1 的 _parameters]
    Sub1 --> SubB1[子 Module 1 的 _buffers]
    Sub1 --> SubM1[子 Module 1 的 _modules]
    style P fill:#dbeafe,stroke:#3b82f6
    style B fill:#fef3c7,stroke:#f59e0b
    style Mods fill:#dcfce7,stroke:#22c55e

9.3 __setattr__：拦截赋值的核心魔法

flowchart TD
    Set["self.x = value"]
    Set --> Q1{value 是 Parameter?}
    Q1 -->|是| P1["注册到 _parameters dict<br/>从 _modules / _buffers 清理"]
    Q1 -->|否| Q2{value 是 Module?}
    Q2 -->|是| P2["注册到 _modules dict<br/>建立子模块树"]
    Q2 -->|否| Q3{value 是 Buffer?}
    Q3 -->|是| P3["注册到 _buffers dict<br/>跟随 .cuda 等迁移"]
    Q3 -->|否| Q4{name 在某个 dict 里?}
    Q4 -->|是| P4["已注册的参数/子模块名只能赋 None<br/>buffer 名可赋 Tensor，否则 TypeError"]
    Q4 -->|否| P5["走 object.__setattr__<br/>普通 Python 属性"]
    style P1 fill:#fef3c7
    style P2 fill:#dcfce7
    style P3 fill:#dbeafe
    style P5 fill:#f3e8ff

torch/nn/modules/module.py:1972 的 __setattr__ 是 nn.Module 最重要的方法，简化版：

def __setattr__(self, name: str, value):
    params = self.__dict__.get("_parameters")

    if isinstance(value, Parameter):
        # 1. value 是 Parameter → 注册到 _parameters
        remove_from(self.__dict__, self._buffers, self._modules, ...)
        self.register_parameter(name, value)

    elif params is not None and name in params:
        # 2. 已存在的 param 被赋 None → 删除
        if value is not None:
            raise TypeError("cannot assign ... as parameter ... ")
        self.register_parameter(name, value)

    else:
        modules = self.__dict__.get("_modules")
        if isinstance(value, Module):
            # 3. value 是 Module → 注册到 _modules
            remove_from(self.__dict__, self._parameters, self._buffers, ...)
            modules[name] = value

        elif modules is not None and name in modules:
            # 4. 已存在的子模块名只能赋 None → 删除；赋别的类型 TypeError
            if value is not None:
                raise TypeError("cannot assign ... as child module ... ")
            modules[name] = value

        else:
            buffers = self.__dict__.get("_buffers")
            if isinstance(value, Buffer) or (buffers is not None and name in buffers):
                # 5. value 是 Buffer 或已存在的 buffer 名 → 注册到 _buffers
                ...

            else:
                # 6. 普通属性 → 走默认 __setattr__
                super().__setattr__(name, value)

——这就是"self.linear1 = nn.Linear(...) 自动注册"的真相。Python 的 __setattr__ 协议被拦截，根据 value 类型路由到不同容器。

9.3.1 为什么要拦截赋值：备选方案的对比

PyTorch 的 __setattr__ 拦截不是唯一方案。其他框架处理"自动注册参数"的不同思路：

• Keras（早期）：要求用户在 build() 方法里显式调 self.add_weight(...)。冗长但显式
• JAX/Flax：用 dataclass 字段声明，外部 framework 通过 reflection 收集字段。需要继承 flax.linen.Module（子类会被自动转成 dataclass）
• TensorFlow：用 tf.Variable 在某个 tf.Module 子类上自动追踪
• PaddlePaddle：与 PyTorch 类似，__setattr__ 拦截

PyTorch 选择 __setattr__ 拦截的好处是 用户写代码与普通 Python 完全一致 —— 不需要装饰器、不需要 build 方法、不需要 add_weight 调用。代价是 nn.Module 是个"魔法重"的类，新人学起来需要一段时间适应它的行为不是普通 Python 类。这种"直观但魔法"vs"显式但冗长"是框架设计的永恒权衡。

9.3.2 一个具体追踪：self.linear = nn.Linear(10, 20)

在 MyModel.__init__ 里写 self.linear = nn.Linear(10, 20) 时：

1. Python 解释器调 MyModel.__setattr__('linear', nn.Linear(10, 20))
2. MyModel.__setattr__ 实际是继承的 Module.__setattr__
3. 检查 value 类型：isinstance(nn.Linear(...), Parameter) → False，isinstance(nn.Linear(...), Module) → True
4. 走分支 3：self._modules['linear'] = nn.Linear(...)
5. 注意 self.linear 这个名字根本没进 self.__dict__

这就是为什么后面 model.linear 还能取到值 —— 通过 __getattr__ 的兜底机制。

注意 __setattr__ 里的 remove_from(...) 调用 —— 它从其他容器里删除同名条目。这是为了避免"同一个名字同时在 _parameters 和 _modules 里"的不一致状态。如果你先 self.x = nn.Linear(...) 再 self.x = nn.Parameter(...)，Linear 会从 _modules 删除、Parameter 进入 _parameters。但反过来不行：先 self.x = nn.Parameter(...) 再 self.x = nn.Linear(...)，会先命中"已存在的参数名"分支（分支 2），直接抛 TypeError: cannot assign 'torch.nn.modules.linear.Linear' as parameter 'x' (torch.nn.Parameter or None expected)（CPU torch 2.11.0 实跑）。已注册的参数名、子模块名只接受同类值或 None，这一点和普通 Python 类的 self.x = ... 并不一样。

9.3.3 一个反直觉的细节：覆盖式赋值如何工作

class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.x = nn.Linear(10, 20)   # 进 _modules
        self.x = nn.Parameter(torch.randn(5))   # 此时 _modules 里的旧值被删, 进 _parameters

第二次赋值时 __setattr__ 检测 value 是 Parameter，先调 remove_from(self.__dict__, self._buffers, self._modules, ...) 删除其他容器里的同名旧值，再 register_parameter。这套"先清后加"语义让这个方向的覆盖式赋值表现得和普通 Python 对象一致 —— 用户感受不到内部容器切换（反方向会报错，见上一节）。

9.3.4 为什么 Parameter 要单独继承

torch.nn.Parameter 的核心作用确实是当作 isinstance 检查的标记，让 __setattr__ 能区分"这是要训练的张量"和"这是普通中间张量"。但它不是一个空壳子类 —— 打开 torch/nn/parameter.py 会看到三处实打实的设计：

class _ParameterMeta(torch._C._TensorMeta):
    # Make `isinstance(t, Parameter)` return True for custom tensor instances that have the _is_param flag.
    def __instancecheck__(self, instance) -> bool:
        if self is Parameter:
            if isinstance(instance, torch.Tensor) and getattr(instance, "_is_param", False):
                return True
        return super().__instancecheck__(instance)


class Parameter(torch.Tensor, metaclass=_ParameterMeta):
    def __new__(cls, data=None, requires_grad=True):
        if data is None:
            data = torch.empty(0)
        if type(data) is torch.Tensor or type(data) is Parameter:
            return torch.Tensor._make_subclass(cls, data, requires_grad)
        # Path for custom tensors: set a flag on the instance to indicate parameter-ness.
        t = data.detach().requires_grad_(requires_grad)
        ...
        t._is_param = True
        return t

三个要点：一，requires_grad 默认就是 True（普通 torch.Tensor 默认 False）—— 这才是"包成 Parameter 就会被训练"的直接原因；二，对普通张量走 _make_subclass 真正换类型，对自定义张量子类（走 __torch_dispatch__ 那一套的）则只在实例上打一个 _is_param 标记、保留原类型；三，_ParameterMeta.__instancecheck__ 让带了这个标记的实例在 isinstance(x, Parameter) 下也返回 True —— 这样 Module.__setattr__ 的类型判断对两条路径都成立。此外 Parameter 还重写了 __deepcopy__ / __repr__ / __reduce_ex__（后者管 pickle 序列化）。

如果你写：

self.weight = torch.randn(10, 20, requires_grad=True)   # ❌ 不会进 _parameters
self.weight = nn.Parameter(torch.randn(10, 20))         # ✅ 进 _parameters

第一种写法 __setattr__ 检测不到 Parameter 类型，直接进 __dict__，不会被 model.parameters() 列出，optimizer 找不到它。这是新手常踩的坑。

9.4 __getattr__：从三个 dict 兜底

torch/nn/modules/module.py:1955 的 __getattr__ 处理"读"路径：

def __getattr__(self, name: str):
    if "_parameters" in self.__dict__:
        _parameters = self.__dict__["_parameters"]
        if name in _parameters:
            return _parameters[name]
    if "_buffers" in self.__dict__:
        _buffers = self.__dict__["_buffers"]
        if name in _buffers:
            return _buffers[name]
    if "_modules" in self.__dict__:
        modules = self.__dict__["_modules"]
        if name in modules:
            return modules[name]
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

Python 的 __getattr__ 只在标准属性查找失败时才被调。所以普通属性走 __dict__ 查找正常路径，特殊属性（参数/buffer/子模块）走兜底。

linear.weight 的查找流程：

1. Python 先查 linear.__dict__['weight'] —— 没有（因为 __setattr__ 没把它放这）
2. 查 type(linear).__dict__['weight'] —— 没有（不是类属性）
3. 走 __getattr__('weight') —— 在 _parameters['weight'] 里找到，返回

这种"读 / 写不对称"是 Python 元编程的经典套路：写靠 __setattr__ 路由到正确容器，读靠 __getattr__ 从容器兜底取出。用户全程无感。

9.4.1 性能取舍：每次属性访问都过 __getattr__ 慢吗

新手经常担心：每次 self.linear 都过 __getattr__，会不会比普通属性访问慢？

结论上确实更慢，但要看清楚慢在哪：__getattr__ 只在标准属性查找失败之后才被调用，所以代价是"一次失败的常规查找 + 一次 Python 层函数调用 + 最多三次 dict 查找"。这比一次命中 __dict__ 的属性访问贵，但绝对值很小 —— 相对一次 forward 里的算子调用可以忽略。真正会被放大的是"每秒上百万次属性访问"的场景，比如反复递归遍历模块树。

这里并没有什么缓存 trick —— 常见的 forward 里每个子模块本来就只访问一次：

def forward(self, x):
    x = self.linear1(x)   # 第一次访问 linear1, 走 __getattr__
    x = self.linear2(x)   # 第二次访问 linear2, 也走 __getattr__
    return x

每个名字只查一次。如果你写 for _ in range(1000): x = self.linear(x)，每次循环都过 __getattr__ —— 这种代码在 PyTorch 里少见但确实有性能损失。torch.compile 在编译时会把这种重复 attribute access 特殊化到具体的子模块对象，消除 __getattr__ 开销。第 12 章 Dynamo 章会展开。

9.5 __call__：模型调用的真实入口

很多人以为 model(x) 等价于 model.forward(x)。错。看 torch/nn/modules/module.py:1918：

__call__: Callable[..., Any] = _wrapped_call_impl

_wrapped_call_impl 本身很短（torch/nn/modules/module.py:1775），它先看有没有编译产物，没有就转给 _call_impl：

def _wrapped_call_impl(self, *args, **kwargs):
    if self._compiled_call_impl is not None:
        return self._compiled_call_impl(*args, **kwargs)
    else:
        return self._call_impl(*args, **kwargs)

真正干活的 _call_impl（torch/nn/modules/module.py:1783）才是那层包裹：它先选 forward_call（trace 中用 _slow_forward，否则用 self.forward），然后是一条很重要的快路径：

# If we don't have any hooks, we want to skip the rest of the logic in
# this function, and just call forward.
if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
        or _global_backward_pre_hooks or _global_backward_hooks
        or _global_forward_hooks or _global_forward_pre_hooks):
    return forward_call(*args, **kwargs)

——一个 hook 都没注册（包括全局 hook）时，module(x) 直接调 forward，不付任何额外代价。只有确实注册了 hook，才会进后面那一大段：跑 forward_pre_hooks（可改 args）→ 需要时用 BackwardHook 包装输入 → 调 forward → 跑 forward_hooks（可改 result）→ 给输出的 grad_fn 挂 non-full backward hook → 异常路径上补跑 always_called 的那些 hook。

所以 model(x) ≠ model.forward(x)。两者的差别在于前者跑 hooks（以及编译产物），后者不跑。这就是为什么生产代码要用 model(x) 而不是 model.forward(x) —— hooks 是 PyTorch 各种工具（torch/nn/utils/prune.py 的剪枝、eager 模式量化的 observer、torch/utils/module_tracker.py 这类模块级追踪等）的扩展点，绕过它们会导致这些工具失效。

9.5.1 _compiled_call_impl 与 torch.compile 的协作

module.py 还有一个重要变量 _compiled_call_impl（类属性，默认 None）。用户调 model.compile()（Module.compile，module.py 的最后一个方法）时，它执行 self._compiled_call_impl = torch.compile(self._call_impl, *args, **kwargs)，此后 __call__ 优先调用 _compiled_call_impl 而不是 _call_impl。

这种"插槽式"设计让 torch.compile 可以就地加速 module，对象本身不变：

model = MyModel()
model.compile()     # 就地填上 _compiled_call_impl，model 仍是 MyModel
out = model(x)      # 走编译后的 _call_impl

注意另一种常见写法 model = torch.compile(model) 走的不是这个插槽：它返回一个包了原模型的 OptimizedModule（torch/_dynamo/eval_frame.py），原模型的 _compiled_call_impl 仍是 None（CPU torch 2.11.0 实跑确认）。

编译本身怎么做见第 12–15 章。注意 __getstate__ 在序列化时会剔除 _compiled_call_impl（因为它依赖运行时编译产物，不能跨进程传递）。这种"序列化时剥离运行时态"是 Module 兼容多种使用场景的细节。

9.5.2 forward 是 Python 函数还是 Module 方法

很多人没意识到：forward 是普通 Python 方法，不是任何特殊声明。Module 的 __call__ 通过 forward_call = self.forward 拿到这个方法，调用时传递参数。这意味着：

• 你可以在 forward 里写任何 Python 代码（if/for/while/try）
• forward 可以接受任何参数（位置 / 关键字 / 默认值）
• forward 可以返回任何东西（tensor、tuple、dict、list、None）

这种自由度是 PyTorch 动态图的体现。代价是：forward 必须在每次 module(x) 都被解释执行（除非用 torch.compile），无法像静态图框架那样跨调用复用编译产物。

子类也可以重写 __call__ 而非 forward —— 但绝大多数代码遵循 "只重写 forward" 的惯例，让 hook 等机制不被破坏。

9.5.3 reentrant 的 forward 调用

如果在 forward 里再调一次 self(x)（递归调用同一个 module），会怎样？

class RecursiveModule(nn.Module):
    def forward(self, x, depth):
        if depth == 0:
            return x
        return self(x, depth - 1)   # 递归调用自己

这是合法的，每一层递归调用都会完整走一遍 _call_impl、各跑一次 hooks（上例 depth=2 时 forward hook 被调 3 次，CPU torch 2.11.0 实跑）。反向图同样按调用展开：每次调用里的算子都会生成各自的 grad_fn 节点，不会复用同一个实例；若递归中反复用到同一个参数，它的 AccumulateGrad 会把多路梯度累加起来 —— 这和 RNN 按时间步重复使用同一组权重是一回事，PyTorch 完全支持。

sequenceDiagram
    autonumber
    participant U as user: model(x)
    participant W as _call_impl
    participant Pre as forward_pre_hooks
    participant F as forward(x)
    participant BH as BackwardHook
    participant Post as forward_hooks
    U->>W: __call__(x)
    W->>Pre: 跑所有 pre-hook
    Pre-->>W: 可能修改 args
    W->>BH: setup_input_hook (注册反向 hook)
    W->>F: forward(args)
    F-->>W: result
    W->>BH: setup_output_hook
    W->>Post: 跑所有 post-hook
    Post-->>W: 可能修改 result
    W-->>U: result

9.6 state_dict：递归收集状态

torch/nn/modules/module.py:2195 的 state_dict 是模型保存的核心。简化逻辑：

def state_dict(self, *, destination=None, prefix="", keep_vars=False):
    if destination is None:
        destination = OrderedDict()
        destination._metadata = OrderedDict()

    local_metadata = dict(version=self._version)
    if hasattr(destination, "_metadata"):
        destination._metadata[prefix[:-1]] = local_metadata

    for hook in self._state_dict_pre_hooks.values():
        hook(self, prefix, keep_vars)

    self._save_to_state_dict(destination, prefix, keep_vars)

    # 递归子模块
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(
                destination=destination,
                prefix=prefix + name + ".",
                keep_vars=keep_vars,
            )

    # 最后跑 state_dict post-hook（_state_dict_hooks，见 §9.6.5 / §9.8）
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, local_metadata)
        ...   # 私有 API 注册的可返回新 dict，公开 API 注册的必须返回 None

    return destination

_save_to_state_dict 把当前 module 的 _parameters 和 persistent _buffers 写到 destination：

# torch/nn/modules/module.py:2144（源码返回类型标注为 None）
def _save_to_state_dict(self, destination, prefix, keep_vars):
    for name, param in self._parameters.items():
        if param is not None:
            destination[prefix + name] = param if keep_vars else param.detach()
    for name, buf in self._buffers.items():
        if buf is not None and name not in self._non_persistent_buffers_set:
            destination[prefix + name] = buf if keep_vars else buf.detach()
    # 省略：子类重写了 get_extra_state() 时，再写入 prefix + "_extra_state"

注意 递归 + prefix 是 dotted path 的来源：

• 调用 model.state_dict()，prefix = ""
• 进 model.linear1.state_dict(prefix="linear1.")
• 在 linear1 里 _save_to_state_dict 把 _parameters['weight'] 存到 destination['linear1.weight']
• 类似产生 'linear1.bias'、'linear2.weight'、'bn.running_mean' ...

这种"前缀 + 递归"的设计让任意深度的 Module 树扁平化为一个 dict[str, Tensor]，序列化非常友好。

值得注意一个工程细节：state_dict 返回的是 OrderedDict + _metadata 字段。_metadata 是个隐藏字段（用 _metadata 而不是 metadata，避免冲突），存每个模块的"版本号"（类属性 Module._version，默认 1）。基类的 _load_from_state_dict 只把它原样透传；需要迁移的子类在自己的 _load_from_state_dict 里读 local_metadata.get("version") 做兼容 —— 例如 BatchNorm 的 _version = 2，加载 version < 2 的旧 ckpt 时会补上当年还没有的 num_batches_tracked（torch/nn/modules/batchnorm.py:109）。这样某个 module 类改了参数 / buffer 布局后仍能加载旧 ckpt。这套版本号机制大多数用户感受不到，但是 PyTorch 跨版本 ckpt 兼容的关键。

9.6.1 forward 与 no_grad 的细微交互

如果你在 forward 里写 with torch.no_grad(): 包住一段代码，会怎样？

def forward(self, x):
    with torch.no_grad():
        prepared = self.preprocess(x)   # 这段不参与反向
    return self.main(prepared)

注意 prepared 没有 grad_fn、requires_grad=False —— 它从 no_grad 块出来。self.main(prepared) 反向时反向链就在 prepared 这里断了：main 的参数照常拿到梯度，但 prepared 不需要梯度，反向到此为止，不会传给 preprocess 的参数。

这是"前处理 / 特征提取不参与训练"的常见写法之一，冻结预训练 backbone、只训练 head 的场景可以这样做；另一种做法是对 backbone 调 requires_grad_(False)（或逐个设 param.requires_grad = False）。

9.6.2 一个常见误解：state_dict 是不是模型本身

新手以为 torch.save(model.state_dict()) 等价于"保存了整个模型"。错：state_dict 只保存权重和 buffer 值，不保存模型类定义。要重建模型必须先 model = MyModel(...) 用 Python 构造，再 model.load_state_dict(...) 灌入权重。

如果想"完整保存模型对象"，要 torch.save(model)（不传 state_dict）—— 这会用 pickle 序列化整个 model 对象。但这种方式有个大坑：pickle 依赖类定义。pickle 只按"模块路径 + 类名"记下类的引用，加载时必须能 import 到同名的类：类被改名、挪了位置就加载失败；类的属性结构改了，旧对象的状态又可能和新代码对不上。另外自 v2.6 起 torch.load 默认 weights_only=True，直接加载整对象会抛 UnpicklingError: Weights only load failed...，必须显式传 weights_only=False（且只对可信来源这样做，CPU torch 2.11.0 实跑）。

所以生产代码强烈建议用 state_dict 而非整对象 pickle：让模型架构（Python 代码）与权重（state_dict）完全解耦，各自独立演进。HuggingFace、PyTorch Lightning 等主流框架都遵循这个原则。

9.6.3 keep_vars=True 的妙用

默认 keep_vars=False：state_dict 里的张量都被 .detach() —— 与参数共享存储，但 requires_grad=False、不再是 Parameter，拿去读写不会牵扯 autograd。

keep_vars=True 则跳过 detach，直接放进原来的 Parameter / buffer 对象本身（实跑：类型是 Parameter、requires_grad=True）。需要按名字拿到"就是模型里那个参数对象"的场景才用得上，比如按 key 把参数和 optimizer 里持有的对象一一对应：

# 值是模型里的 Parameter 对象本身，而不是 detach 后的视图
params_by_name = dict(model.state_dict(keep_vars=True))

普通用户用默认 detach 就够。

9.6.4 _non_persistent_buffers_set 的设计

不是所有 buffer 都该进 state_dict。比如 BatchNorm 的 num_batches_tracked 是必要的（要存），但某些用户自己加的"调试用 running stats"可能不该存。Module.register_buffer(name, tensor, persistent=False) 让用户标记某 buffer 不进 state_dict —— 名字加入 _non_persistent_buffers_set，_save_to_state_dict 时跳过。

这种"buffer 是否持久化"的细分是 v1.6.0 才加入的（v1.5.0 的 module.py 里还没有 _non_persistent_buffers_set）。

9.6.5 自定义 state_dict 输出格式

某些场景需要修改 state_dict 的 key（如要转成另一套库约定的命名）。可以注册 state_dict post-hook：

def rename_keys(module, state_dict, prefix, local_metadata):
    # 把 module.weight 改成 module.W
    if prefix + 'weight' in state_dict:
        state_dict[prefix + 'W'] = state_dict.pop(prefix + 'weight')

model.register_state_dict_post_hook(rename_keys)   # 公开 API，要求就地修改 state_dict

对一个 nn.Linear 注册后，state_dict() 的 key 从 ['weight', 'bias'] 变成 ['bias', 'W']（CPU torch 2.11.0 实跑）。在不同命名约定之间转换 checkpoint 时，这是 nn.Module 层面现成的扩展点。

9.7 load_state_dict：参数加载与名字匹配

torch/nn/modules/module.py:2531 的 load_state_dict(state_dict, strict=True, assign=False) 反向操作。它做：

1. 递归遍历模块树（先父后子），每层 module 调 _load_from_state_dict：先跑本模块的 _load_state_dict_pre_hooks，再把 state_dict 里对应 prefix 的 key 拷到自己的 _parameters / 持久 _buffers
2. 每层 module 在子模块都加载完之后，跑本模块的 _load_state_dict_post_hooks
3. 收集 missing_keys / unexpected_keys / error_msgs：strict=True 时把前两者也并入错误信息；只要 error_msgs 非空就 raise RuntimeError，否则返回 (missing_keys, unexpected_keys)。不存在只 warn 的分支

最常见的"加载报错"场景：训练时模型结构变了，旧 ckpt 的 key 与新模型对不上。strict=False 让加载继续但记录差异，方便迁移学习场景。但要注意它只放过"key 对不上"，不放过形状不符：fc 层名字没变、只是 num_classes 从 1000 改成 10 时，直接加载仍会抛 RuntimeError: ... size mismatch for fc.weight: copying a param with shape torch.Size([1000, ...]) from checkpoint, the shape in current model is torch.Size([10, ...]).。正确做法是先把旧 fc 的 key 去掉：

# 加载预训练 ResNet, 替换最后的 fc 层
model = MyResNet(num_classes=10)
state_dict = torch.load('resnet50.pth')
state_dict = {k: v for k, v in state_dict.items() if not k.startswith('fc.')}
missing, unexpected = model.load_state_dict(state_dict, strict=False)
# missing = ['fc.weight', 'fc.bias']（新 fc 保持随机初始化），unexpected = []

第 19 章序列化章会展开 torch.save / torch.load 与 state_dict 的协作，包括 safetensors 格式的优势。

9.7.1 missing_keys 与 unexpected_keys 的诊断价值

load_state_dict 返回的 missing_keys / unexpected_keys 是诊断"模型 vs ckpt 不匹配"的金矿：

• missing：当前模型有但 ckpt 没有的 key —— 通常是模型新加了层，ckpt 是老版本
• unexpected：ckpt 有但当前模型没有的 key —— 通常是模型删了层，或 ckpt 多保存了不必要的 buffer

调试技巧：

missing, unexpected = model.load_state_dict(state_dict, strict=False)
print('Missing keys:', missing[:5])
print('Unexpected keys:', unexpected[:5])

差异只有零星几项时，常见原因是 BN 的 num_batches_tracked 这类次要 buffer 或少数新增 / 删除的层；大面积对不上，通常意味着模型架构有重大变化（或者 key 前缀整体错位，比如多了一层 module.），加载后行为可能不正常。

9.7.2 加载流程的细节

load_state_dict 实际工作流：

flowchart TB
    Start[load_state_dict 入口]
    Start --> Load[递归 load:<br/>先父后子, 对每个 module 调 _load_from_state_dict]
    Load --> Each["_load_from_state_dict 做:<br/>1. 跑本模块的 _load_state_dict_pre_hooks<br/>2. 把 state_dict 里 prefix+name 的 tensor 拷到 _parameters/_buffers<br/>3. 记录 missing/unexpected keys 与 error_msgs"]
    Each --> Post["子模块都加载完后<br/>跑本模块的 _load_state_dict_post_hooks"]
    Post --> Check{"strict=True 且有 missing/unexpected,<br/>或有 error_msgs(如 size mismatch)?"}
    Check -->|是| Err[raise RuntimeError]
    Check -->|否| Return[返回 missing_keys, unexpected_keys]
    style Each fill:#fef3c7,stroke:#f59e0b

注意默认 数据拷贝是 inplace：state_dict 里的 tensor 在 torch.no_grad() 下用 param.copy_(input_param) 拷到现有的 Parameter，不是创建新对象。这样原模型的 Parameter 引用保持有效（optimizer 的 param_groups 不会失效），但权重值已经更新。这是为什么 load 后不需要重新 build optimizer。传 assign=True 时则相反：用 setattr 把 state_dict 里的张量（必要时包成 Parameter，requires_grad 沿用模块原值）直接换进模块，参数对象变了（实跑：load_state_dict(..., assign=True) 之后 w is model.weight 为 False），所以源码文档要求此时 optimizer 在加载之后再建；它常配合 meta device 上构造的模型使用，省掉一次拷贝。

这种 "in-place loading" 设计有个隐形约束：state_dict 里的 tensor 必须形状与现有 Parameter 完全一致。形状不一致会报 size mismatch，而且 strict=False 也不会放过它。如果你想换形状（如改 num_classes），要先重新构造 model，或者像上面那样先从 state_dict 里删掉那几项再加载。

9.8 9 类 hook：模块生命周期的扩展点

nn.Module 提供 9 类 hook（Module.__init__ 里为它们准备了 11 个 OrderedDict：8 个存 hook 本身，其中新旧两种 backward hook 共用 _backward_hooks；另 3 个 _forward_hooks_with_kwargs / _forward_hooks_always_called / _forward_pre_hooks_with_kwargs 只是按 hook id 记标记、当集合用），按调用时机分类：

Hook	注册方法	触发时机
forward_pre_hook	register_forward_pre_hook	forward 调用前，可改 args
forward_hook	register_forward_hook	forward 调用后，可改 result
full_backward_pre_hook	register_full_backward_pre_hook	反向开始前
full_backward_hook	register_full_backward_hook	反向结束后，能拿 grad_input / grad_output
backward_hook (deprecated)	register_backward_hook	早期版本，与 full_backward 不兼容
state_dict_pre_hook	register_state_dict_pre_hook（torch/nn/modules/module.py:2131）	state_dict 调用前
state_dict_post_hook	register_state_dict_post_hook（torch/nn/modules/module.py:2107）；老的私有 _register_state_dict_hook（torch/nn/modules/module.py:2088）仍在	state_dict 收集完之后
load_state_dict_pre_hook	register_load_state_dict_pre_hook（torch/nn/modules/module.py:2305）	load_state_dict 之前
load_state_dict_post_hook	register_load_state_dict_post_hook（torch/nn/modules/module.py:2317）	load_state_dict 之后

两个 state_dict post-hook API 的区别值得单独记一笔：私有的 _register_state_dict_hook 允许 hook 返回一个新的 state_dict，但这个返回值只对根 module 生效、对子 module 会被丢掉（这是 issue #117437 记录的 bug）；公开的 register_state_dict_post_hook 修掉了这个不一致 —— 它只允许 hook 就地修改 state_dict。源码用挂在 hook 函数上的 _from_public_api 标记来区分两者，并禁止同一个函数同时走两条注册路径。

每种 hook 都是为特定扩展场景设计的：

• profiler / 性能分析：用 forward_pre/post hook 记时间戳
• 混合精度 cast：用 forward_pre hook 把 fp32 输入转 fp16
• 梯度统计 / 调试：用 full_backward_hook 检查 grad_input / grad_output
• 量化 observer：eager 模式量化用 forward_hook 收集激活分布（torch/ao/quantization/quantize.py 的 _observer_forward_hook）

（DDP 的 gradient bucketing 也是"某个 grad 算完就触发 AllReduce"，但它挂的是参数 AccumulateGrad 上的 autograd hook，不是 Module hook，见第 17 章。）

• 检查点平均：state_dict_hook 在 save 时拦截

每条 hook 都是 PyTorch 与生态对接的关键点。理解 hook 机制能让你优雅扩展任何 nn.Module 行为，不需要改源码。

9.8.1 hook 与 dispatcher mode 的关系

第 5 章 §5.7 我们看到 TorchDispatchMode 也是一种"拦截每个算子"的扩展点。Module hook 与它的区别：

维度	Module hook	TorchDispatchMode
拦截粒度	一整个 Module 的 forward	每一个 ATen 算子
注册方式	module.register_forward_hook(fn)	with MyMode():
看到的输入	Module 的输入参数	ATen 算子的张量参数
适合场景	"在某层之后做 X"	"对所有算子做 X"

两者不是互斥的 —— 你可以同时用。Module hook 是粗粒度（按层），DispatchMode 是细粒度（按算子）。生产代码里大部分扩展用 Module hook（更易理解），需要看到每个算子时才用 DispatchMode（如 FakeTensorMode 的假执行）。

9.8.2 RemovableHandle 设计

每个 register_xxx_hook 都返回 RemovableHandle：

handle = model.register_forward_hook(my_hook)
# 后续...
handle.remove()    # 移除

RemovableHandle（torch/utils/hooks.py）是个很简单的普通类：构造时从全局自增计数器 next_id 取一个 id，并用 weakref.ref 弱引用注册它的字典（外加可选的 extra_dict，例如 with_kwargs 那几个标记字典）。remove() 就是从这些字典里删掉这个 id。

为什么不直接让用户把 hook 函数传回来撤销？因为同一个函数可能被注册多次（在同一个或不同 hook 集上），handle 的 id 让"第 N 次注册"可识别可撤销。RxJava 的 Disposable 也是这种"注册返回句柄、凭句柄撤销"的思路。

9.8.3 一个高级技巧：借 __setattr__ 做 model surgery

"model surgery" 是指在不改源码的前提下修改训练好的模型行为。例如：把 ResNet 的某层 ReLU 换成 GELU，但只改一层、保持其他层原样。

def replace_relu_with_gelu(module):
    for name, child in module.named_children():
        if isinstance(child, nn.ReLU):
            setattr(module, name, nn.GELU())
        else:
            replace_relu_with_gelu(child)

replace_relu_with_gelu(model)

这段代码递归遍历所有 module，遇到 ReLU 就替换成 GELU。setattr(module, name, ...) 触发 __setattr__，自动从 _modules 删除旧的 ReLU、加入新的 GELU。整套 surgery 操作零侵入、不需要改原模型代码。

这种能力让 PyTorch 在研究场景里非常灵活 —— 可以快速尝试"如果某层换成 X 会怎样"。HuggingFace 的 PEFT 库（LoRA 等）就大量用 model surgery 把"普通 Linear 替换成 LoRALinear"。

9.8.4 一个真实例子：grad clipping with hook

def grad_clip_hook(module, grad_input, grad_output):
    return tuple(g.clamp(-1.0, 1.0) if g is not None else g for g in grad_input)

model.register_full_backward_hook(grad_clip_hook)

要看清它裁的是什么：full backward hook 返回的新 grad_input 是对模块输入的梯度，只影响继续往上游（更早的层）传的梯度；本模块自己参数的 .grad 不受影响。实跑一个两层 nn.Sequential 挂上这个 hook：输入 x.grad 被压到 1.0 以内，而第一层 weight.grad 的最大绝对值是 104（CPU torch 2.11.0）。所以它不能代替 torch.nn.utils.clip_grad_norm_ 这类对参数梯度的裁剪，生产代码裁梯度仍用标准 API。

9.8.5 hook 执行的隐形顺序

每种 hook 之间的调用顺序是有确定 spec 的，按 _call_impl 实现的顺序：

1. global forward_pre_hook（用 nn.modules.module.register_module_forward_pre_hook 注册，对所有 module 生效）
2. self.forward_pre_hook（per-module）
3. 若有 full backward hook 或 backward pre-hook，用 BackwardHook 包装输入（setup_input_hook）
4. self.forward(args)
5. global forward_hook
6. self.forward_hook
7. backward_hook 设置（包装输出）
8. non_full_backward_hook 注册到 grad_fn

理解这套顺序很重要。比如你写一个 forward_hook 想看到"经过其他所有 hook 修改之后的最终 result"，要确保你的 hook 注册时机最晚（dict 是按插入顺序）；反过来想最先跑，就用 prepend=True 参数。注意 prepend 只在本模块的 hook 之间调序，global hook 永远排在前面（实跑顺序：g_pre → pre0(prepend) → pre1 → forward → g_post → post0(prepend) → post1）。这种顺序敏感性是 Module hook 与 dispatch mode 的另一个区别 —— DispatchMode 用栈结构、最近的 mode 先拦截，逻辑更清晰。

9.8.6 用 hook 实现 GPT-style cache：一个高级例子

class KVCache:
    def __init__(self):
        self.cache = {}

    def hook(self, module, input, output):
        # 假设 module 是 attention 层
        layer_id = id(module)
        self.cache[layer_id] = output[1]   # K, V

cache = KVCache()
for layer in model.transformer.layers:
    layer.attention.register_forward_hook(cache.hook)

这是 transformer 推理时手动截取 K/V 的 toy 写法：forward hook 只能把每层 attention 的输出存下来，要在下一次 forward 时真正复用，还得改 attention 的 forward 去读这份缓存。生产代码的 KV cache（如 vLLM 的 PagedAttention）并不靠 hook，而是在 attention 内部直接读写分页缓存。这个例子演示的是 Module hook 让"不改模型代码就把中间结果截出来"非常直接。

9.9 train() / eval()：模式切换

def train(self, mode: bool = True):
    if not isinstance(mode, bool):
        raise ValueError("training mode is expected to be boolean")
    self.training = mode
    for module in self.children():
        module.train(mode)
    return self

这个方法把 self.training 设为 True/False 并递归子模块（eval() 就是 self.train(False)）。BatchNorm、Dropout 等模块在 forward 时依据 self.training 决定行为。下面是 Dropout 行为的示意（真实源码只有一行 return F.dropout(input, self.p, self.training, self.inplace)，由 F.dropout 在 training=False 时原样返回输入）：

class Dropout(Module):
    def forward(self, input):
        if self.training:
            return F.dropout(input, self.p, ...)
        else:
            return input    # eval 模式下 dropout 是 identity

新手最常见的 bug 就是 inference 时忘了 model.eval() —— BN 仍然用 batch 统计、Dropout 仍然丢神经元、模型行为非确定。生产 inference 服务必须 model.eval()。

注意 eval() 和 no_grad() 是正交的概念：

• eval()：切到推理模式（影响 BN/Dropout 行为）
• no_grad()：关掉 autograd（影响梯度计算）

生产 inference 通常两者都要：

model.eval()
with torch.no_grad():
    out = model(x)

或者更高效的：

model.eval()
with torch.inference_mode():
    out = model(x)

9.9.1 model.children() vs model.modules()

两个递归方法容易混：

• children() 只返回直接子模块（不递归到子模块的子模块）
• modules() 递归返回所有模块（包括 self、子子模块……）

类似还有 named_children() 和 named_modules() 带名字版本。生产代码里：

# 给所有 Conv2d 的 weight 做 init
for m in model.modules():
    if isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight)

modules() 按深度优先、先自身后子模块的顺序遍历整个网络，并且默认去重：它建立在 named_modules(remove_duplicate=True) 上，用 memo 集合记下已产出的实例，同一个子模块被挂在多处（如 nn.Sequential(l, l)）也只产出一次。这是 PyTorch 给"按类型批量处理参数"留的最方便接口。

children() 与 modules() 看似冗余，但场景不同：

• 想给"每个直接子层"应用变换（如最外层 wrap 一下）：用 children()
• 想搜整个网络里所有某类型层（如所有 BN）：用 modules()
• 想知道某层在网络里的命名位置：用 named_modules()

apply(fn) 并不走 modules()：它的实现是 for module in self.children(): module.apply(fn)，然后 fn(self)、return self —— 先子后父的后序遍历，顺序与 modules() 相反，也不做全树去重（实跑 nn.Sequential(Linear, Sequential(ReLU))：apply 的访问顺序是 Linear → ReLU → 内层 Sequential → 外层 Sequential，modules() 则是外层 Sequential → Linear → 内层 Sequential → ReLU）。这种"递归遍历 + 用户提供闭包"模式让 PyTorch 与训练流程框架（如 Lightning）能轻松实现"对所有参数做 X"的统一接口。

9.9.2 train/eval 在分布式与 inference 服务中的微妙性

model.train() / model.eval() 看似简单，分布式训练里有些坑：

1. DDP 包装后的 train/eval：DDP 把 model 包成 DistributedDataParallel(model)，调 ddp_model.train() 会递归设置 inner module 的 training。但有些第三方库的 wrapper 不递归，需要显式 ddp_model.module.train()

2. 模型分阶段切换：典型场景是先训练后用一段 inference 子图。常见错误：

model.eval()
out = model.encoder(x)   # eval 模式
loss = some_loss(out)
loss.backward()           # 仍然会反向, 但 BN 用的是 running stats 而非 batch stats
                          # 训练效果会很差

正确做法：明确分开训练阶段和推理阶段，不要混合。

3. 推理服务的并发安全：model.eval() 设置了 self.training=False，多线程并发推理时不会互相影响（因为每个推理只读 training）。但如果你某条路径里又调了 model.train()，并发的其他推理会被污染。生产服务里不要在请求处理路径切换 train/eval。

9.10 跨专栏关联

• 《vLLM 推理内核深度解析》专栏第 7 章（模型加载）：vLLM 加载 HuggingFace 模型时直接构建 PyTorch nn.Module 树，再用自定义 weight loader 喂权重 —— 它不走 load_state_dict，而是按本章的 dotted 参数名把权重逐个映射、切分后拷进参数，依赖的是参数注册与命名机制
• 《Tokio 源码深度解析》专栏第 6 章（Task 的生命周期）：Tokio Task 的 spawn / hook 注册机制与 nn.Module 的 hook 思想类似，都是"在统一抽象上预留扩展点"
• 《Serde 元编程》专栏第 12-13 章（派生宏代码生成）：Serde 的 #[derive(Serialize)] 在编译期生成的逻辑，与 nn.Module 在运行期通过 __setattr__ 拦截做的事是同一目的：自动收集类的字段
• 《MCP 协议设计与实现》专栏第 5 章（Tool 注册）：MCP server 用 decorator 注册 method handler 与 nn.Module hook 注册思想极相似 —— 都是"对类的字段做声明式处理"

9.11 一个对照：torch.fx vs nn.Module

最后值得提一下 torch.fx。它是 PyTorch 用来"把 nn.Module 转成符号化 IR"的工具：

import torch.fx

model = MyModel()
gm = torch.fx.symbolic_trace(model)
# gm 是一个 torch.fx.GraphModule, 它的 forward 是一段被记录的 IR
print(gm.graph)

fx.symbolic_trace 用 Proxy 对象代替真实输入跑一遍 forward，借 __torch_function__ 协议和运算符重载把每次调用记录成 IR 节点；nn.Linear、nn.BatchNorm1d 这类标准层默认当作叶子，整层记成一个 call_module 节点（上例实跑打印出 placeholder → call_module[target=linear1] → call_module[target=bn] → call_module[target=linear2] → output）。FX 在 v1.8.0 引入，比 torch.compile（v2.0）早；如今 Dynamo 取代的是 symbolic_trace 这个抓图前端，FX Graph 本身仍是 Dynamo / AOTAutograd / Inductor 之间传递的 IR。

理解 fx 与 nn.Module 的关系：nn.Module 是面向人的 API（layers + forward），fx Graph 是面向编译器的 IR。两套表示互相可转换：fx.symbolic_trace 把 Module 转成 Graph；fx.GraphModule 又是一种 Module（继承自 nn.Module）。这种"原始 + 编译"两套表示并存是 PyTorch 框架的经典哲学，也是第 12-15 章 torch.compile 章节的根本设计。

注意抓图这件事现在由 Dynamo 主导（它更鲁棒、能处理更多 Python 控制流），但 symbolic_trace 在某些特定场景仍然有用：

• 量化 (quantization)：FX 图模式量化（torch/ao/quantization/quantize_fx.py）的 prepare / convert 都在 FX 图上做；PT2E 量化同样操作 FX 图，只是 v2.11 已从主仓迁往 torchao
• 模型变换 / pruning：剪枝、重参数化等需要明确的 Graph IR 才好做
• 简单的 trace 场景：用户想看"我的 Module 实际跑了哪些算子"时，symbolic_trace 比 Dynamo 简单得多

第 12 章 Dynamo 章会展开 fx 与 Dynamo 之间的接力（Dynamo 产出的正是 FX Graph）。

9.12 几条工程建议

实战 nn.Module 相关的最佳实践：

1. 永远 super().__init__() 在 __init__ 第一行：忘了会让 _parameters 等字典没初始化，后续赋值子模块报 AttributeError: cannot assign module before Module.__init__() call（赋 Parameter 时是 cannot assign parameters before Module.__init__() call）

2. 学习参数用 nn.Parameter 包：直接 self.weight = torch.randn(...) 不会被自动注册为参数

3. 不参与训练但要保存的状态用 register_buffer：如 BN 的 running stats、自定义的位置编码

4. 临时中间状态不要注册成 buffer：register_buffer / nn.Buffer 注册的持久 buffer 会被 state_dict 收进来，让 ckpt 变大。用局部变量，或者 persistent=False。普通 self.x = torch.zeros(...) 赋的张量属性（名字不是已注册的 buffer 时）只进 __dict__，不进 state_dict，但也不会跟着 .to() / .cuda() 搬迁（实跑：to(torch.float64) 后 buffer 变成 float64，普通张量属性仍是 float32）；纯 Python 对象（list、dict、自定义类）同样不会被收集 —— 只有 Parameter、持久 buffer（外加子模块里的这两者）才进 state_dict

5. 子模块要写到 __init__ 而非 forward：在 forward 里 nn.Linear() 每次都新建模块，权重不会被训练

6. 注意 hook 的 order 影响：多个 hook 按注册顺序调用，写的时候要明确"我能假设其他 hook 已经跑过吗"

7. model.to(device) 是 inplace 操作：返回 self，但同时也改了 self。new = model.to(device) 与 model.to(device) 等价

8. nn.ModuleList / nn.ModuleDict 而非 Python list / dict：self.layers = [nn.Linear(10,10) for _ in range(5)] 不会自动注册（普通 list 不被 __setattr__ 检测），self.layers = nn.ModuleList([...]) 才会

9. nn.Parameter 注册时机：要在构造 optimizer 之前注册，一般就在 __init__ 里。运行期动态加 Parameter（如 if cond: self.weight2 = nn.Parameter(...)）会让 optimizer 拿不到（因为 optimizer 构造时只读了一次 parameters()，之后得手动 add_param_group）

10. model.apply(fn) 递归调用：apply 给每个 module 调一次 fn(module)。常用于初始化：model.apply(init_weights)。它的实现是先对每个 children() 递归 apply、再 fn(self)、最后 return self（先子后父，见 §9.9.1）

11. forward 里不要用 self.training 之外的 mode 切换：除了 BN / Dropout 这种内置约定，自定义 module 不要在 forward 里读其他 self.xxx 切换行为。Dynamo 会把这类属性值特化进图、装上 guard，值一变就重新编译（实跑：model.compile(backend="eager") 后把这样一个 bool 属性来回切，产生 2 张图、没有 graph break），切换频繁时编译开销会反复出现

12. 模型结构包含可变 list / Tensor list：参数列表用 nn.ParameterList（torch/nn/modules/container.py:646），例如 self.scales = nn.ParameterList([nn.Parameter(torch.ones(1)) for _ in range(K)])，自动注册。注意 PyTorch 没有 nn.BufferList（v2.11 的 torch/nn/modules/container.py 里的容器类是 Sequential / ModuleList / ModuleDict / ParameterList / ParameterDict，外加已弃用的 Container，没有 buffer 版本）—— 一组 buffer 只能逐个 self.register_buffer(f"buf_{i}", t)，或者塞进一个专门的子 Module 里

9.13 给 Module 设计师的"通用启示"

把 nn.Module 思想抽象到任何"组件容器"系统：

第一：__setattr__ 拦截 + __getattr__ 兜底是 Python 元编程的黄金组合 —— 让用户感受不到框架在背后做的事，写代码像写普通类一样自然

第二：子节点用专门容器（dict / list）而非 self.x = child：让递归遍历、序列化、迁移操作都有统一入口

第三：hook 系统让框架可扩展：在生命周期关键点开放可注册的回调列表，第三方工具就能集成进来。RemovableHandle 让用户能撤销注册

第四：state 与 metadata 分离：state_dict 存"值"，_metadata 存"版本号 / 类型信息"。前者让数据可移植，后者让兼容性成为可能

第五：train/eval 是模型生命周期的两个不同阶段：用一个 self.training flag 切换，相关组件查这个 flag 调整行为。这种"全局上下文 + 组件就地查询"模式比"为不同模式写两份代码"清晰得多

第六：让递归遍历成为一等接口：parameters() / modules() / children() / named_xxx() 这种"按容器维度遍历"接口让框架的所有变换、保存、迁移操作都有标准入口

第七：train/eval / persistence / device 这三种"全局开关"互不干扰：每种开关由不同字段表示，Module 在合适的层面响应。这种"正交概念分离"让用户能任意组合而不互相破坏

把这七条记住，你写自己的"组件框架"（如 GUI 框架、游戏引擎的 Entity-Component 系统、CRM 的工作流模板）能少走很多弯路。

9.14 nn.Module 的调用开销由什么构成

nn.Module 在每次 module(x) 上加的都是 Python 层开销，源码里能一条条数清楚：

成本项	何时发生	源码位置
_wrapped_call_impl → _call_impl 两层 Python 函数调用	每次调用	torch/nn/modules/module.py:1775 / 1783
一次 torch._C._get_tracing_state() 判断选 forward_call	每次调用	_call_impl 开头
8 个 hook 容器的布尔判空	每次调用	_call_impl 的快路径 if
pre/post hook 的逐个 Python 调用、BackwardHook 包装	仅当注册了 hook	快路径之后的分支
__getattr__ 兜底查三个 dict	每次访问 self.某子模块 / 某参数	torch/nn/modules/module.py:1955

三条定性结论比背数字有用得多。第一，没注册任何 hook 时快路径直接 return forward_call(*args, **kwargs)，额外成本只剩两层函数调用加一串判空 —— 这也是为什么"随手给全模型挂个 global forward hook"是有代价的：它会让所有 module 掉出快路径。第二，这些开销按 module 调用次数计，与张量多大无关；所以层数多、每层计算小的模型（小 batch、小 hidden、CPU 上的小网络）里它的占比才会显现，大模型的大 GEMM 面前完全淹没。第三，torch.compile 之所以在这一侧有收益，是因为 Dynamo 把整段 Python（含属性查找与调用开销）一起 trace 掉了，剩下的只有 kernel。

要判断自己的模型属于哪一类，用 torch.profiler 看 CPU 侧的 Python 时间与 device 侧的 kernel 时间比例，别照搬任何现成数字。

9.15 nn.Module 的 to/cuda/half 实现

model.cuda() 把整个模型搬到 GPU。它的实现就是用 _apply 递归遍历每个 parameter / buffer，对每个张量调 .cuda()：

def cuda(self, device=None):
    return self._apply(lambda t: t.cuda(device))

def _apply(self, fn, recurse=True):
    if recurse:
        for module in self.children():
            module._apply(fn)
    # apply to all parameters
    for key, param in self._parameters.items():
        if param is not None:
            with torch.no_grad():
                param_applied = fn(param)
            self._parameters[key] = Parameter(param_applied, ...)
    # apply to all buffers
    for key, buf in self._buffers.items():
        if buf is not None:
            self._buffers[key] = fn(buf)
    return self

_apply 是 nn.Module 里最通用的一个抽象 —— 它接受任意"对张量的变换"函数，递归整个网络。.cuda() / .cpu() / .half() / .float() / .bfloat16() / .to(device, dtype) 全部用它实现。

上面是简化版。真实的 _apply 对参数分了三条路（buffer 则简单得多，一句 self._buffers[key] = fn(buf) 直接覆盖）：

• torch.utils.swap_tensors(param, param_applied) —— 当开了 torch.__future__.get_swap_module_params_on_conversion()，或者变换结果 param_applied 是 traceable wrapper subclass、原参数 param 是 FakeTensor 时走这条，原地把两个张量的内容对调，param 这个 Python 对象本身不变
• param.data = param_applied —— 当 torch._has_compatible_shallow_copy_type(param, param_applied) 成立（新旧张量类型兼容）且没开 get_overwrite_module_params_on_conversion() 时走这条，同样保留原对象
• self._parameters[key] = Parameter(param_applied, param.requires_grad) —— 兜底路径，换一个新的 Parameter 对象

为什么要有前两条？因为 .cuda() 这类变换会返回新张量，如果每次都换新对象，那些已经持有旧 Parameter 引用的地方（最典型的就是 optimizer 的 param_groups）就会指向失效的对象。CPU 上实跑 model.to(torch.float64)，走的是 .data = 这条路，model.linear1.weight 前后是同一个对象。不过 Module.cuda() 的文档仍写着 "This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized" —— 兜底路径（以及打开 overwrite 标志时）确实会换对象，所以"先搬 device，再建 optimizer"仍是稳妥的顺序。

9.15.1 _apply 与 autograd 的微妙交互

_apply 对参数调 fn 时用 with torch.no_grad(): 包住，避免变换本身被 autograd 记录。源码注释说得清楚：模块里的参数是图的叶子，不希望 param_applied 带上 autograd 历史 —— 否则 .cuda() 得到的新张量会挂着 _to_copy 的 grad_fn，不再是叶子。

这种"对底层 housekeeping 操作显式禁用 autograd"是 PyTorch 内部很多基础设施的共同做法。optimizer.step() 内部也用 no_grad 包，避免参数更新被反向追踪。第 10 章会再看到这种模式。

9.15.2 _apply 的边界：直接转 int8 行不通

def quantize_to_int8(t):
    return t.to(torch.int8) if t.is_floating_point() else t

quantized = model._apply(quantize_to_int8)
# RuntimeError: data set to a tensor that requires gradients must be floating point or complex dtype

看上去是"一行量化"，实跑（CPU torch 2.11.0）会在 param.data = param_applied 那一步报上面的错：参数 requires_grad=True，而整数张量不能要求梯度。_apply 适合 .half() / .bfloat16() / .to(device) 这类"逐张量换 dtype / device"的变换；真正的量化要把模块换成量化模块、把权重打包成量化格式，AMP autocast 则是在 dispatcher 层按算子改 dtype，都不是站在 _apply 上写的。第 20 章量化会展开。

9.15.3 _save_to_state_dict vs _load_from_state_dict 的对偶性

为什么 PyTorch 把 save 和 load 写成 _save_to_state_dict / _load_from_state_dict 这种"_私有方法"模式而不是公开 API？

因为 它们是 module 子类可以重写的扩展点。如果你的 module 有特殊持久化需求（如某个张量需要量化后存、或者把多个张量打包成一个），可以重写这两个方法：

class MyModule(nn.Module):
    def _save_to_state_dict(self, destination, prefix, keep_vars):
        # 自定义存储格式
        destination[prefix + 'packed'] = self._pack_for_save()

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, ...):
        # 自定义加载逻辑
        if prefix + 'packed' in state_dict:
            self._unpack(state_dict[prefix + 'packed'])

这种"框架提供 hook 点 + 子类重写细节"是 Template Method 设计模式的应用。torch/ao/nn/quantized/ 下的量化模块（含量化 Embedding）、torch/ao/nn/sparse/ 的稀疏量化 Linear，以及 nn.BatchNorm* / nn.InstanceNorm*，都重写了 _load_from_state_dict 来处理特殊情况。

9.16 一个完整的 Module 树：HuggingFace BERT 解构

实战看一段 HuggingFace BertModel 的 model.named_modules() 输出（节选）：

'' (BertModel)
'embeddings' (BertEmbeddings)
'embeddings.word_embeddings' (Embedding)
'embeddings.position_embeddings' (Embedding)
'embeddings.token_type_embeddings' (Embedding)
'embeddings.LayerNorm' (LayerNorm)
'encoder' (BertEncoder)
'encoder.layer.0' (BertLayer)
'encoder.layer.0.attention.self.query' (Linear)
'encoder.layer.0.attention.self.key' (Linear)
'encoder.layer.0.attention.self.value' (Linear)
'encoder.layer.0.attention.output.dense' (Linear)
'encoder.layer.0.intermediate.dense' (Linear)
'encoder.layer.0.output.dense' (Linear)
'encoder.layer.1' (BertLayer)
...

每条路径就是 dotted prefix，对应 state_dict 里的 key。这种"树结构 + dotted path"的设计让 BERT 这种 12 层 / 24 层的复杂模型能用扁平 dict 完整表达，加载、迁移、部署都极简单。HuggingFace、Lightning 等所有上层框架都建立在 nn.Module 这套机制之上。

观察一个细节：encoder.layer.0 这种带数字的路径来自 nn.ModuleList。它在构造、append / extend、__setitem__ 时都把列表索引转成字符串名字（'0'、'1' 等），经 add_module / setattr 存到 _modules 字典。所以 ModuleList 表面上是列表，底层仍然是 dict —— 索引化路径自然形成。

类似地，nn.ModuleDict 让用户用字符串 key 索引子模块，但它的所有 key 都直接进 _modules。理解 ModuleList / ModuleDict / Sequential 都只是 dict 的语法糖，能让你写自己的"结构化 module 容器"时心里有数。

nn.Sequential 与 ModuleList 是兄弟而不是父子 —— torch/nn/modules/container.py 里三者都直接写作 class Sequential(Module)、class ModuleList(Module)、class ModuleDict(Module)。Sequential 的特别之处只在于它自己实现了 forward：按 _modules 的插入顺序遍历子模块，逐个 input = module(input)。常用于线性堆叠的网络：

mlp = nn.Sequential(
    nn.Linear(10, 100),
    nn.ReLU(),
    nn.Linear(100, 10)
)

mlp[0] 是第一个 Linear、mlp[-1] 是最后一个 Linear。这是 nn.Module 给"简单堆叠模型"的便利包装。但 Sequential 不能处理需要分支 / 跳连接的复杂结构，那些场景要写完整的 forward。

9.17 一个收官的反思：nn.Module 的成功是 PyTorch 的成功

某种程度上，PyTorch 之所以击败 TF 1.x 赢得学术界，不只是动态图，还因为 nn.Module 比 tf.keras.Model 更自然。

tf.keras.Model 早期要求用户走 functional API（x = Dense(10)(x)）或者 subclass + build() 方法 —— 都比 PyTorch 的 self.linear = nn.Linear(...) 多了一层心智负担。研究者要快速实验新架构时，PyTorch 的 nn.Module 写起来更接近"在草稿纸上画的网络结构图"。

这条经验给所有想做"用户友好框架"的设计师一个启示：让用户写代码的方式贴近他思考的方式。如果用户脑子里想的是"线性层 → ReLU → 线性层"，框架就该让他这样写代码。中间多一个抽象层都是认知摩擦。

PyTorch 的 nn.Module 把"魔法量"控制在 两个 dunder 方法 + 几个内部 dict 这个最小集 —— 已经能让用户体验最自然，又不至于黑盒到完全不可调试。这种 "隐藏复杂度但保留可观察性" 的工程平衡，是 nn.Module 设计上最高明的一笔。

下一章拆 optimizer：optim.SGD / optim.Adam 等怎么实现，包括最近引入的 Foreach / Fused / Capturable 三种性能模式。