老一批用 TensorFlow 1.x 写过模型的人应该还记得那种“给每个变量手动取名、还要四处收集参数列表”的日子。后来切到 PyTorch几乎所有模型都长成一个样子class MyNet(nn.Module)。torch.nn.Module 是 PyTorch 里绕不开的基座类你用它的频率可能比用nn.Linear还高但很多人对它的理解只停留在“一个用来放层的容器”。说实话这个理解不算错可一旦后面涉及到参数冻结、子模块复用、多卡训练、模型导出各种奇奇怪怪的报错就会冒出来。这篇内容我想从一个写过不少 PyTorch 代码的人的角度把 torch.nn.Module 的参数追踪机制、子模块树结构、设备搬运逻辑以及那些只有实际踩过坑才会注意到的细节一次讲透希望能帮助真正准备长期使用 PyTorch 的人把地基打好。1. 从“一个类”到“一切模型的基座”nn.Module 到底为你做了什么一个刚接触 PyTorch 的人通常会先照抄一段最简单的模型代码import torch.nn as nn class MLP(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): return self.fc2(torch.relu(self.fc1(x)))这段代码看起来平平无奇好像只要把“想用的层”放到__init__里再在forward里把它们串起来就算完事了。但如果你一直带着这种理解写模型后面大概率会在某个时刻被卡住为什么我用列表保存的子层训练了半天参数一个都没动为什么model.cuda()之后某个张量还在 CPU 上为什么明明冻结了参数优化器还在更新它这些问题的答案全部指向同一个点nn.Module不是一个简单的“层容器”它是一套状态和结构的管理框架负责把参数、子模块、缓冲区、设备信息全部串起来。1.1 没有 nn.Module 的模型会是什么样子为了理解 nn.Module 的价值先想象一下没有它的日子。你可以自己定义一个类在__init__里创建一堆张量当权重再写一个forward方法做矩阵乘法和激活函数。这样模型确实能跑。但接下来你要面对一系列“手工活”想拿到全部参数做梯度裁剪必须自己遍历所有层。想把模型搬到 GPU 上每个张量都要手动.to(cuda)。想保存模型得自己维护一个字典确保每个权重名字不冲突。想实现类似 PyTorch 的model.train()/model.eval()状态切换你得给每个需要状态控制的层都写开关。我早年试过一次仅为一个 10 层左右的 CNN 实现上面这几件事就写了两百多行样板代码而且每加一个层都要小心翼翼去更新索引逻辑稍微漏一个参数模型就默默把某个权重留在 CPU 上。后来我才意识到nn.Module 把这一整套公用的管理逻辑都提取出来了你需要做的只是继承它然后照着规则把组件放进去。1.2 nn.Module 是用“树形结构”换取“自动化管理”神经网络本质上是一棵“有向计算图”中间节点是各种子模块叶子节点是参数输入从根节点进入沿着结构一路算到输出。nn.Module 做的事情就是把这棵树完整记录下来并给你一套递归操作它的工具。举个最直观的例子下面是同样一件事有 nn.Module 和没有 nn.Module 的对比需求没有 nn.Module使用 nn.Module获取全部参数手工遍历所有层自己汇总model.parameters()一行搞定搬到 GPU逐个张量.cuda()model.to(cuda)递归处理冻结网络自己标记每个张量module.requires_grad_(False)保存权重自己维护格式state_dict()自动给出有序状态嵌套组合自己维护父子关系赋值子模块后自动注册所以你表面上写的是“一个类”实际上你是在维护一棵递归可遍历、自带状态管理的参数树。理解了这一点后面很多行为就顺理成章了。2. 参数、子模块、缓冲区nn.Module 内部那“三本账”的追踪逻辑我第一次接触 nn.Module 时产生过一个疑问为什么在__init__里写self.weight nn.Parameter(torch.randn(...))这个 weight 就能自动被model.parameters()抓到是不是 PyTorch 在背后做了某种“魔法”答案确实是魔法但这个魔法其实很朴素nn.Module重写了__setattr__方法。每当你执行self.xxx yyy的时候它都会先拦截住然后根据yyy的类型决定把记录写进哪一本“账”。2.1 三本账的注册规则Parameter、Module、Buffer在nn.Module.__init__()执行之后每个模块实例都会维护三个有序字典_parameters、_modules、_buffers。它们各自的职责是_parameters存放nn.Parameter类型的可训练参数。_modules存放子nn.Module形成树结构。_buffers存放非训练但需要随模型保存的状态比如 BatchNorm 的 running_mean、running_var。当你在__init__里这样赋值时背后的行为是不同的self.weight nn.Parameter(torch.randn(4, 4)) # 自动进入 _parameters self.fc nn.Linear(4, 4) # 自动进入 _modules self.register_buffer(running_mean, torch.zeros(4)) # 注册进 _buffers注意一个细节普通 Tensor 直接赋值给一个属性不会自动注册成参数也不会注册成 buffer。很多人一开始会写self.running_mean torch.zeros(4)想当然地认为模型就能跟踪它结果发现保存权重时根本没有这个值。只有在_buffers里已经注册过的名字后续给它赋值才能覆盖到缓冲中。如果你希望某个张量既不算可训练参数又要进 state_dict就必须用register_buffer显式注册。2.2 参数和缓冲区在模型生命周期中的差别搞清楚 Parameter 和 Buffer 的区别能避免很多潜在问题。二者虽然都会随着model.to()一起移动也会进入state_dict()但在语义上完全不同行为nn.ParameterBuffer进入model.parameters()会不会默认需要梯度默认requires_gradTrue默认requires_gradFalse进入model.state_dict()会会优化器默认更新会不会典型用途权重、偏置BN 统计量、位置编码缓存这也是很多人写自定义模块时容易搞混的地方。比如实现一个类似 LayerNorm 的模块想把里面的 scale 当Parameter把计算过程中的均值缓存当 Buffer如果类型用反了训练时就会莫名其妙多出一些“参数”被优化器更新或者某些参数梯度永远为 0。2.3 用 List[nn.Module] 存子模块为什么会“失联”这一条我觉得是最值得反复强调的坑。网上不少老代码里会这样写class BadNet(nn.Module): def __init__(self): super().__init__() self.layers [nn.Linear(8, 8) for _ in range(3)] # 错误写法表面上看self.layers里有三个 Linearforward 里也能正常调用self.layers[0](x)。但实际上这三层里的权重根本不会被注册。因为nn.Module.__setattr__只认nn.Module类型的对象而一个 Python list 不是 Module它内部的元素不会递归进入账本。结果就是model.parameters()返回空或者缺参数model.to(cuda)把这几个 Linear 留在 CPU优化器也不会更新它们。报错往往很晚才出现比如 forward 时提示设备不一致或者训练 loss 一直不变。正确做法是使用nn.ModuleListself.layers nn.ModuleList([nn.Linear(8, 8) for _ in range(3)])同样如果你用的是 Python dict 装子模块就要替换成nn.ModuleDict。这条规则没有任何例外只要你想让一个“容器里的模块们”参与模型自身的参数管理它就必须是nn.ModuleList或者nn.ModuleDict。3. 手写自定义模块时你其实在写“结构与行为的契约”很多时候我们并不需要自定义模块直接用nn.Sequential堆叠现成层就够了。但一旦你开始写类似 Transformer Block、自定义注意力、带特殊初始化策略的网络就得自己手写一个继承 nn.Module 的类。这时候很多人会把注意力全放在 forward 里怎么写计算却忽略了继承 nn.Module 的真正意义你写的__init__是在声明结构forward是在声明行为而 nn.Module 负责把结构和行为绑定在一起。3.1 为什么是继承而不是简单组合你可能会想“我只要把层当作普通成员变量做一个普通类再写一个 forward 方法不也行吗”从纯数学计算角度确实行但你就失去了前面说的那套自动化能力。继承 nn.Module 的好处是模型的参数自动进入model.parameters()不用手工收集。模型可以参与嵌套被父模块当作子模块注册。model.train()、model.eval()、model.to()这些操作会自动作用于内部组件。PyTorch 生态里的torch.compile、FX 图优化、Hook 机制都默认以 nn.Module 作为单位。还有一个更实际的原因如果你在普通类里保存了 nn.Module 层这些层自己仍然需要一个父级账本。与其临时用register_module挂一个假父节点不如老老实实继承 nn.Module让一切在语义上都是通的。3.2 forward 只是被call“包装”的普通方法新手经常问我定义了 forward调用时为什么要写model(x)而不是model.forward(x)因为__call__方法在真正调用 forward 前后还会执行一系列 hook 逻辑。大致流程是执行model(x)时PyTorch 会进入模块的__call__运行已注册的 forward-pre hook然后调用你的forward最后再运行 forward-post hook。这个机制也被用到梯度缓存、剪枝、蒸馏、模型导出等场景。如果有一天你为了绕开某些逻辑直接调用model.forward(x)你就把这些 hook 链完全跳过了行为可能和预期不一致。所以一个基本规则是对外永远用model(x)不要直接调self.fc1.forward(x)即使在你的类内部也别这么干。3.3 一个带初始化策略的完整自定义模块演示一个自定义的“线性 LayerNorm Dropout ReLU”模块并给它挂上初始化策略import torch import torch.nn as nn class MLPBlock(nn.Module): def __init__(self, in_dim, hidden_dim, dropout0.1): super().__init__() self.linear nn.Linear(in_dim, hidden_dim) self.norm nn.LayerNorm(hidden_dim) self.dropout nn.Dropout(dropout) def forward(self, x): x self.linear(x) x self.norm(x) x torch.relu(x) return self.dropout(x)这里__init__的作用只是“声明我有哪些组件”不在其中做任何计算。权重初始化可以放在外面用model.apply递归遍历所有子模块def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) nn.init.zeros_(m.bias) block MLPBlock(16, 32) block.apply(init_weights)apply这个函数非常实用它会按照树的遍历顺序把每个子模块都访问一遍。理解了这条你也就知道怎么批量改模块属性了。3.4 模块可以被多次调用但要注意状态性同一个 nn.Module 实例在一个 forward 里可以被调用多次这在很多结构里都会发生。线性层这样纯函数式的组件没问题调用两次就是两次独立计算。但 Dropout、BatchNorm 这类带状态或随机性的层在同一个 forward 中被多次调用时行为会叠加。比如 Dropout 在训练模式下每次 forward 都会生成新的 mask如果你在 forward 里调用了同一个包含 Dropout 的模块两次两次的 dropout mask 互不相同这通常是预期行为但如果你天真地把结果缓存就会出错。自定义模块时要清楚自己挂载的子模块是否是“有状态”的以及它是否受 train/eval 模式影响。4. 嵌套与复用模块拼装中的所有权、共享与梯度写真实项目时几乎不可能只用一层嵌套。一个标准模型往往是若干子模块组合Encoder 里套若干 BlockBlock 里套 Attention、FFN、Norm。nn.Module 对这种嵌套是天然支持的你在__init__里赋值的任何子模块都会被自动注册到父模块的_modules字典里形成一棵树。4.1 所有权模型与打印结构当父模块“持有”一个子模块时子模块的参数就属于父模块这棵树的叶子节点。model.parameters()会递归收集所有叶子节点的参数。也正是因为这种递归结构你可以很方便地用print(model)看到层级每多一层嵌套就会多一级缩进。如果你希望快速查看某个网络每一层的参数量和输出形状可以用torchinfo.summary(model, input_size(batch, ...))它本质上是向模型丢一个假输入然后在内部遍历 nn.Module 树统计出每个子模块的参数量。这个工具能在早期帮你发现“哪层参数数量不对”、“哪层没有注册成功”。4.2 同一个模块对象被多处引用时的共享问题我见过不少项目为了节省参数会让多个子模块共享同一个权重对象。例如编码器和解码器共享同一个 Embedding。下面这样写是常见做法self.embed nn.Embedding(1000, 256) self.encoder EncoderBlock(self.embed) self.decoder DecoderBlock(self.embed)如果EncoderBlock和DecoderBlock在__init__里把传入的embed赋值成自己的属性比如self.shared_embed embed那么同一个 Embedding 对象就会同时出现在两个子模块的_modules里。这会导致几个连锁现象model.parameters()里会包含同一个weight参数对应的两个引用。model.state_dict()中会出现两个不同的 key但值对应同一个存储。优化器在初始化参数组时会提示某参数在参数组中出现多次实际更新时会造成重复更新。这里不是说共享权重一定会出大错而是说你要清楚这样做之后PyTorch 内部对参数的管理不再是一一对应的关系。我的建议是如果确实要共享最好让共享权重的对象只被注册一次子模块内部不要重复注册它而是把已经算好的嵌入向量通过 forward 参数传入或者把共享模块放在父级子模块只接收经过它处理后的结果。4.3 梯度回传时的叠加与重复更新即使你绕开了重复注册只要同一个 Parameter 参与了 forward 里的多条计算路径反向传播也会把各条路径的梯度累加到该参数上。这是符合 autograd 自然行为的也是共享权重模型的标准语义。真正的坑往往出现在“参数重复注册 优化器不去重”的组合。举个实际场景你把同一个 Module 赋给了两个不同名字的父模块属性然后直接用model.parameters()构造优化器。第一次 step 时这个参数的梯度会被累计优化器会因为同一个张量对象出现多次而更新两次最后权重变化幅度直接翻倍。虽然新版 PyTorch 会在初始化优化器时给出 UserWarning但很多人没认真看警告就把它忽略了。排查这种问题最直接的方法是打印一下list(model.parameters())的 id 列表看看有没有重复对象。5. 设备、模式与状态搬运.to()、.train/.eval 和 state_dict 的联动逻辑早期用 PyTorch 时我常犯一个错误只对输入张量调用了.cuda()忘了模型本身还在 CPU 上然后 forward 报出设备不匹配。后来才意识到model.to()是一个深度遍历过程但它也遵循某些规则搞清楚这些规则能少走不少弯路。5.1 为什么 model.to(device) 能“传染”到所有子模块因为 nn.Module 内部维护了_parameters和_buffers这两本账model.to()会遍历当前模块以及所有子模块对每个 Parameter 和 Buffer 执行相应类型转换和设备移动。这个过程是递归向下的所以哪怕你的模型嵌了五六层子模块一次调用就能全部搬过去。但要注意普通 Python 属性不会被处理。如果你在__init__里保存了一个self.some_tensor torch.zeros(...)并且它不是通过register_buffer注册的那么model.to()根本不会移动它。最直接的后果是 forward 里一旦用到它设备不匹配又会出现。遇到这种“明明调用过 to() 还是报设备错误”的情况先检查是不是有张量藏在了普通属性里。5.2 train/eval 模式的决定权在谁手上model.train()做的事同样是从根节点开始把树上的每个模块的training标志设为 True。model.eval()则相反。这影响的是对模式敏感的层Dropout 在 training 为 True 时生效BatchNorm 会更新滑动统计量在 eval 模式下Dropout 变成恒等函数BN 使用累计的 running_mean 和 running_var。如果你自己写了一个需要区分 train/eval 行为的模块可以直接在 forward 里读self.training。但如果你重写了train()或eval()方法记得调用super().train(mode)否则子模块的模式不会被正确设置。这个细节很隐蔽我在自定义模块里曾经因为忘了调 super 而让 Dropout 永远处于训练状态一直到最后验证阶段才发现。5.3 state_dict 的保存与加载本质是一次“键值拷贝”model.state_dict()返回一个 OrderedDict包含模型所有参数和缓冲区key 是它们在树中的路径名。模型保存与加载就变成了两个非常对称的操作torch.save(model.state_dict(), model.pt) model.load_state_dict(torch.load(model.pt))load_state_dict默认是strictTrue它要求加载字典的 key 与当前模型的 key 完全匹配多一个少一个都会报错。这个严格模式其实是好事能避免你漏掉某些层。我在实践中会反复打印state_dict().keys()来做结构对账尤其是改模型结构之后。比如把某层从self.fc nn.Linear(...)改成self.fc nn.Sequential(nn.Linear(...))同一个逻辑层在 state_dict 里的 key 就变了旧权重必然加载失败。用strictTrue去加载通常能第一时间发现问题如果图省事改成strictFalse可能会让某些层静默用随机初始化训练曲线会告诉你结果有多糟。6. 冻结参数、状态加载以及我踩过的 nn.Module 隐藏坑这一节我想专门聊聊微调和冻结。现在很多人拿到预训练模型第一件事就是把主干冻结只微调头部。nn.Module给你提供了非常便捷的接口requires_grad_(False)。但“便捷”的背后有一些容易被忽略的行为会异常。6.1 requires_grad_(False) 并不能阻止所有更新先说参数本身。你执行for p in model.parameters(): p.requires_grad_(False)后这些参数确实不会计算梯度了。正常情况下如果继续使用这些参数构建 loss反向传播时它们的.grad会是 None。很多人因此认为“冻结完成了”。但有个细节如果你构造优化器时直接传了model.parameters()而不是用过滤条件filter(lambda p: p.requires_grad, model.parameters())那么优化器依然会把这些 requires_gradFalse 的参数放进参数组。虽然 step 时它们的梯度是 NonePyTorch 会跳过更新但参数组里带着一堆无效参数除了浪费内存还可能在混合精度场景下触发奇怪警告。更隐蔽的是 BatchNorm。即使你把 BN 里的 weight 和 bias 冻结了只要模型处于 train 模式BN 层的 running_mean 和 running_var 依然会不断更新。因为这些统计量是 buffer不属于参数梯度系统。所以如果你希望彻底冻结一个 BN 子模块需要module.eval()配合requires_grad_(False)并注意全局模型是否还处于 train 状态。6.2 共享模块和参数名冲突一个真实复现的加载错乱案例我接过一个项目里面有一个变量为了“省显存”把同一个 LayerNorm 对象同时赋给了self.norm_a和self.norm_b。训练时一切正常因为 forward 里两个路径都调用了同一个对象计算是对的。但保存权重以后state_dict 里同时出现了norm_a.weight和norm_b.weight两个 key它们指向同一个存储。加载时麻烦来了用旧版本保存的权重可能只有norm_a.weight而新结构还需要norm_b.weightstrict 模式就会报 missing key。有人图省事加了strictFalse结果norm_b变成了随机初始化模型行为直接变了。我当时排查时反复对比 key最后才发现是共享模块对象导致的状态分裂。所以我要强调一个原则你能通过某种“引用”让模块在逻辑上共享但最好不要让同一个 Module 对象被多个属性名同时注册。如果实在要共享请只在顶层注册一次其他位置用函数参数或纯 forward 传值的方式接收数据避免 state_dict 出现重复 key。6.3 排查 nn.Module 树结构的有效手段当模型行为不符合预期时我通常会做这样一套排查也分享给大家打印repr(model)确认子模块归属关系是否符合预期。打印list(model.state_dict().keys())检查 key 是否重复、是否少某些层。用torchinfo.summary检查每一层参数量和输出形状定位参数异常的模块。如果怀疑模块没有被注册用list(model.named_modules())看它的完整模块树确认目标层是否真的在这个树里。这些手段配合起来基本能定位绝大多数与 nn.Module 结构相关的问题。尤其是named_modules()它远比直接打印模型全貌更可控因为它返回的是带路径名的模块枚举可以精确匹配某个叶子节点。最后再分享一个小经验我在写自定义模型时会把“自定义参数容器”这件事尽量简化。能用nn.Sequential组合的绝不手写forward必须手写forward的__init__里只放必要组件需要被模型跟踪的额外状态一律走register_buffer。这样一套纪律执行下来我踩到 nn.Module 相关隐藏坑的次数确实少了很多。torch.nn.Module 并不难用难点在于你要理解它背后那套“账本”和“递归遍历”的设计逻辑只要把这个底层概念焊死在脑子里很多报错你一眼就能看出问题出在哪。