torch.nn

torch.nn 以 nn.Module 为核心抽象，把参数、子模块、缓冲区组织成一棵可遍历、可序列化的树；理解注册机制、模式切换与 state_dict 导出，就掌握了整个模块体系。

在 PyTorch 架构总览的分层中，torch.nn 处于网络搭建层：向下消费 Tensor 与 Storage 的张量运算，向上向 torch.optim 输出待优化参数。这一层的全部 API 围绕同一个意象展开——模型即一棵树：树怎么建、挂什么、怎么遍历、放哪些预制构件、行为由哪个开关控制、哪些活外包给函数式接口、最终如何初始化与落盘。

nn.Module 生命周期

torch.nn 的一切都挂在 nn.Module 这一个抽象上：它把模型组织成一棵树，一个 Module 从构造到落盘经历四个阶段——注册 → 模式切换 → 执行 → 导出。四个阶段里，模式切换与导出在后文各有专门小节，本节先看每天都在写的两个动作：注册与执行。

注册发生在 __init__。nn.Module 重写了 __setattr__：属性赋值为 nn.Parameter 时进 _parameters，赋值为 nn.Module 时进 _modules，普通 Tensor 不被追踪（具体原理查看 notes_nn_module_setattr）。自定义模型只需按惯例写属性，框架自动维护整棵树：

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class Net(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.fc1 = nn.Linear(dim, dim * 2)          # 子模块，进 _modules
        self.scale = nn.Parameter(torch.ones(1))    # 可训练参数，进 _parameters
        self.register_buffer("step", torch.zeros(1))  # 非训练状态，进 _buffers
 
    def forward(self, x):
        return F.relu(self.fc1(x)) * self.scale

执行走 __call__ 而非直接调 forward()：__call__ 先触发 forward pre-hooks / hooks，再进入用户实现的 forward。前向过程由 torch.autograd 记录计算图，Module 本身不感知梯度——Module 只管组织与状态，求导完全外包。

三类注册物与遍历 API

__setattr__ 的分流决定了树上的挂载物分三类，类别直接决定一个张量与优化器、序列化的关系——会不会被训练，还是只是随行状态：

注册物	存储	特征	典型例子
`nn.Parameter`	`_parameters`	`requires_grad=True`，随优化器更新	权重、偏置
buffer	`_buffers`	不参与训练，但随 `.to()` 与 `state_dict()` 迁移	BN 的 running_mean/var
子模块	`_modules`	递归汇聚前两者的注册物	层、块

register_buffer(name, tensor, persistent=True) 注册不训练但需随模型保存/迁移的张量；persistent=False 则不进 state_dict，适合纯运行期统计。nn.Parameter 是 Tensor 的薄子类，仅默认开启梯度。

注册是写入，遍历是读出，两者共享同一棵树——优化器靠 parameters() 收齐全部叶子参数，序列化沿树压平出带点号路径的键名，.to() 沿树递归搬运设备。四个遍历入口的差异在广度与去重：

方法	范围	返回
`children()`	直接子模块	不递归
`modules()`	递归所有后代（含自身）	去重
`parameters()`	递归所有参数	可直接喂优化器
`named_*()` 变体	同上	附带 `(name, obj)`

apply(fn) 递归对每个子模块调用 fn，是批量初始化与 requires_grad_(False) 冻结的惯用载体——nn.init 一节会再次用到它。

容器与常用层

树建好之后是往里放东西：网络骨架靠容器搭，具体的层大多不必手写。nn.Sequential 是有序容器，forward 按序执行，适合直线前馈网络；nn.ModuleList 把列表中的层注册进树，按名索引的对应物是 nn.ModuleDict / nn.ParameterList / nn.ParameterDict——容器怎么选、list 陷阱如何排查，展开见 nn.Sequential 与 ModuleList。

普通 Python list 不注册

self.layers = [nn.Linear(4, 8), nn.Linear(8, 2)] 这样写在 __init__ 里不报错，但 __setattr__ 只认 Parameter 与 Module，list 两者都不是——其中的参数不进 _parameters，喂优化器时凭空消失。这是 nn 最经典的静默陷阱。

预制层覆盖了几乎所有常见结构，每类都有专门笔记承接细节：全连接与卷积的形状约定、权重布局分别见 nn.Linear 与 nn.Conv2d；归一化族的参数布局与 running stats 机制见 nn 归一化层；LSTM 的张量约定与隐状态管理见 nn.LSTM；查找表的 padding_idx 与 weight tying 见 nn.Embedding；inverted dropout 及其变体见 nn.Dropout。损失函数同属此层：CrossEntropyLoss 内部即 log_softmax + NLLLoss，输入应为 raw logits，勿再手动 softmax，类不平衡场景的工具见 nn.CrossEntropyLoss。

train() 与 eval()

预制构件之所以省心，一半功劳在这个全局开关：带状态的层（BN、Dropout）的前向行为取决于 self.training 标志。model.train() / model.eval() 递归翻转每个模块的标志，不改变任何参数值，只改变前向行为；分化最典型的是这两类层：

层	training=True	training=False
BatchNorm	用 batch 统计归一化，并更新 running 统计	用 running 统计，停止更新
Dropout	按概率随机置零并按 $1/ (1 - p)$ 缩放	恒等映射

推理忘调 eval()，失真不止来自 Dropout

BN 会被测试 batch 的统计持续污染，running_mean/var 在每次前向都被改写；单样本推理时 batch 方差无意义，输出直接失真。

Module 与 nn.functional 的选型

train/eval 的自动生效有个前提：状态封装在 Module 里。nn.functional（惯用别名 F）是另一个极端——无状态的函数式接口，两者按是否携带状态分工：

带参数或带模式（BN/Dropout）→ 用 Module：参数自动注册，train/eval 自动生效
纯无参操作（F.relu、F.pad、F.softmax）→ 在 forward 里直接调 F.*，不为它们建层
F.conv2d / F.linear 需手动传权重，仅在需要显式权重控制（手写 attention、weight tying）时使用；此时 F.dropout(x, training=self.training) 必须显式传模式

混合写法是主流惯例：__init__ 只建有状态层，forward 内组合 F.*。

nn.init 初始化策略

预制层自带默认初始化，够不够用取决于信号方差：初始化决定方差在前向/反向是否稳定传播。nn.init.*_ 系列均在 no_grad 下原地写入，配合上文的 apply 使用：

def init_weights(m: nn.Module):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, nonlinearity="relu")  # ReLU 族
        if m.bias is not None:
            nn.init.zeros_(m.bias)
 
model.apply(init_weights)

两种主流策略的差别在方差目标，分野是激活函数饱和与否：

策略	方差目标	适用激活
Xavier/Glorot（`xavier_normal_`）	$\frac{2}{f a n _{in} + f a n _{o u t}}$	tanh/sigmoid 等饱和激活
Kaiming/He（`kaiming_normal_`）	$\frac{2}{f a n _{in}}$	ReLU/LeakyReLU 族

ReLU 负半轴置零损失一半方差，故 Kaiming 只除以 $f a n_{in}$ ；mode="fan_out" 改为保反向方差。nn.Linear 默认按 Kaiming uniform（ $a = 5$ ）初始化，多数场景够用；显式初始化留给残差分支、attention 等对方差敏感的结构。

state_dict 与序列化

注册机制的最终兑现是导出：state_dict() 沿树压平，返回参数名到张量的有序字典，含全部参数与 persistent buffer，不含网络结构——结构由代码定义，存档只存权重，因此加载必须先重建模型：

torch.save(model.state_dict(), "net.pt")  # 只存张量字典
 
model = Net(dim=64)  # 代码重建结构
model.load_state_dict(torch.load("net.pt", map_location="cpu"))  # strict=True 默认键须完全匹配

strict=True 要求键集合严格一致；迁移学习部分加载用 strict=False 并核对缺失/意外键。weight tying（如 Embedding 与输出层共享权重）会让同一张量以不同键名出现两次。整模型 pickle、safetensors 及跨框架格式的完整版图见模型权重保存。

知识花园

探索

torch.nn

nn.Module 生命周期

三类注册物与遍历 API

容器与常用层

train() 与 eval()

Module 与 nn.functional 的选型

nn.init 初始化策略

state_dict 与序列化

关系图谱

目录

反向链接