PyTorch 里写模型绕不开nn.Module这个类。很多人刚上手时会有个错觉以为只要把网络层堆进一个类里、继承一下Module前向传播能跑通就算学会了。可真到了要改结构、要冻结部分参数、要复用子模块、要排查“为什么参数没更新”的时候才发现自己对Module的理解只停留在表面。这篇内容就围绕Module这个核心类展开把它到底是什么、内部怎么组织、实际写模型时哪些地方最容易出问题结合我这些年踩过的坑讲清楚。适合已经能跑通简单训练脚本、但想真正把模型组织方式搞明白的读者也适合刚接触 PyTorch、想一开始就养成正确习惯的新手。1. Module 到底在模型里扮演什么角色1.1 它不是一个“容器”而是一套参数与状态的登记系统刚学的时候我习惯把Module理解成一个“装层的盒子”继承它就是为了有个地方写forward。这个理解不算错但太浅了。Module真正的价值在于它是一套登记系统你往类里赋值的每一个子模块、每一个Parameter都会被自动登记到_modules和_parameters这两个字典里。正是有了这套登记后面parameters()、state_dict()、.to(device)、.train()/.eval()这些操作才能一次性作用到整棵模块树上。换句话说forward只是告诉 PyTorch“数据怎么流”而Module的登记机制告诉 PyTorch“这个模型有哪些东西需要被管理”。这两件事是分开的。很多人写模型只关心前者结果在保存、加载、迁移设备、冻结参数这些环节反复出问题根子都在后者没理解透。举个最典型的例子。下面这段代码看起来没问题import torch import torch.nn as nn class MyModel(nn.Module): def __init__(self): super().__init__() self.layers [nn.Linear(10, 10) for _ in range(3)] def forward(self, x): for layer in self.layers: x layer(x) return x它能跑forward也没错但model.parameters()会是空的。原因很简单你把层放进了一个普通的 PythonlistModule的__setattr__拦截不到列表内部的元素这些层根本没被登记。正确写法是用nn.ModuleList或者nn.Sequential。这个坑我见过太多人踩而且往往是在训练跑完发现 loss 不降、回头查半天才定位到。1.2 参数、缓冲、子模块三类被管理对象的区别Module管理的对象大致分三类搞清楚它们的差异很多“玄学问题”就迎刃而解了。类型登记位置是否参与梯度是否进 state_dict典型用途Parameter_parameters是是可学习的权重、偏置buffer_buffers否是默认均值方差、位置编码等常量子Module_modules递归管理是层、块、整个子网络Parameter是Tensor的子类只有被赋值为类属性时才会自动登记。如果你把一个Tensor直接赋给self.w torch.randn(10)它不会成为参数也不会被优化器看到。要让它可学习必须包一层nn.Parameter。buffer是很多人忽略的一类。像 BatchNorm 里的running_mean、running_var就是 buffer它们不需要梯度但需要随模型一起保存和迁移设备。自己写模型时如果有这类“跟着模型走但不训练”的张量用register_buffer注册比直接存成普通属性靠谱得多。我早期写过一个带固定位置编码的模型把编码矩阵存成普通属性结果.to(cuda)之后编码还在 CPU 上前向直接报设备不一致。改成register_buffer就一劳永逸。1.3 forward 与call的分工别在 forward 里做初始化新手常问为什么调用模型是model(x)而不是model.forward(x)因为Module定义了__call__它会在真正执行forward前后插入一系列钩子逻辑比如 forward pre-hook、forward hook。这些钩子在做特征可视化、梯度检查、调试中间输出时非常有用。如果你直接调forward这些钩子全部失效。由此引出一条实操纪律初始化逻辑放__init__数据流逻辑放forward别在forward里创建层或参数。我见过有人在forward里根据输入维度动态建Linear第一次跑没事第二次跑参数就重复注册、设备也对不上。动态结构确实有场景需要但要用nn.ModuleList配合条件判断在__init__阶段规划好而不是在forward里临时new一个层出来。2. 自定义 Module 的标准骨架与常见写法取舍2.1 一个能直接抄的模板以及每行的意图下面这个骨架是我自己写新模块时的起点几乎每个自定义层都可以从它演化出来import torch import torch.nn as nn class CustomBlock(nn.Module): def __init__(self, in_dim, out_dim, dropout0.1): super().__init__() self.norm nn.LayerNorm(in_dim) self.proj nn.Linear(in_dim, out_dim) self.act nn.GELU() self.drop nn.Dropout(dropout) self.register_buffer(scale, torch.tensor(1.0)) def forward(self, x): x self.norm(x) x self.proj(x) x self.act(x) x self.drop(x) return x * self.scale逐行说一下意图。super().__init__()必须第一行调用否则_parameters等字典还没初始化后面赋值子模块会直接报错。子模块按执行顺序声明方便阅读时对照forward。register_buffer那行演示了常量的注册方式scale会跟着模型走但不参与训练。forward里只做张量运算不产生任何新的被管理对象。这里有个细节值得说nn.GELU()这类无参数层其实也可以直接在forward里用F.gelu(x)函数式写法。两者区别在于函数式写法不占_modules条目state_dict里也不会多出空条目。对于无状态的操作函数式更轻对于有参数或有状态的层必须用模块形式。我个人的习惯是有参数、有 buffer、需要切换 train/eval 行为的一律用模块纯函数变换用F里的函数。2.2 Sequential、ModuleList、ModuleDict 该怎么选这三个容器是组织子模块的主力但用途差别很大选错了会带来不必要的限制。nn.Sequential适合线性串联且不需要在中间取值的场景。它的优点是简洁forward都不用自己写。缺点是中间层输出拿不到也没法做残差连接这类需要跳线的结构。如果你发现自己在Sequential外面又套了一层手动forward去取中间结果那说明这个场景不该用Sequential。nn.ModuleList适合需要按索引访问、需要循环遍历、需要动态增删的场景。它只负责登记不负责前向forward要自己写。写多层堆叠的 Transformer block 时ModuleList几乎是标配class Stack(nn.Module): def __init__(self, num_layers, dim): super().__init__() self.blocks nn.ModuleList([ CustomBlock(dim, dim) for _ in range(num_layers) ]) def forward(self, x): for block in self.blocks: x block(x) return xnn.ModuleDict适合按名字索引的场景比如多任务模型里每个任务一个头或者按配置动态选择分支。它的键必须是字符串访问时self.heads[task_a]这样取。三者共同的底线是绝对不要用普通 Python list/dict 装子模块。前面已经说过原因这里再强调一次因为这是最高频的隐性 bug 来源。2.3 参数初始化的正确位置与手法初始化必须放在__init__里且要在子模块声明之后。PyTorch 默认会给Linear、Conv做初始化但很多时候我们需要自定义比如按层类型分别处理def init_weights(m): if isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.LayerNorm): nn.init.ones_(m.weight) nn.init.zeros_(m.bias) class Model(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(128, 256) self.norm nn.LayerNorm(256) self.fc2 nn.Linear(256, 10) self.apply(init_weights)self.apply(fn)会递归遍历所有子模块并调用fn是批量初始化的标准手法。注意apply是在所有子模块都声明完之后调用顺序不能反。另外如果你在__init__里手动对某个参数做了赋值要确认它没有被后续的apply覆盖或者干脆把特殊初始化放在apply之后单独处理。有个容易忽略的点nn.init的函数默认是原地操作返回的也是同一个张量但如果你写self.fc1.weight nn.init.trunc_normal_(self.fc1.weight)这种赋值形式虽然结果一样但语义上容易让人误以为创建了新张量。直接用原地调用即可不必重新赋值。3. 参数管理冻结、分组、设备迁移的实操细节3.1 requires_grad 的批量控制与常见误操作冻结部分参数做微调是高频需求。标准做法是遍历named_parameters()按名字匹配后设置requires_gradfor name, param in model.named_parameters(): if name.startswith(backbone): param.requires_grad False这里有个坑冻结之后要把模型切到 eval 模式或者至少确认 BatchNorm/Dropout 的行为符合预期。requires_gradFalse只影响梯度计算不影响train()模式下 BatchNorm 更新 running stats。如果你冻结 backbone 是为了做线性探测那 backbone 里的 BN 层应该处于 eval 状态否则 running stats 还在变等于没完全冻住。这一点在迁移学习里特别关键我见过有人冻结了参数但 loss 曲线诡异最后发现是 BN 在偷偷更新。另一个误操作是冻结后忘记过滤优化器。如果你在冻结之前就建好了优化器那优化器里仍然持有这些参数的引用。虽然requires_gradFalse的参数不会产生梯度优化器更新时因为梯度是None通常也不会动它们但更稳妥的做法是先冻结、再建优化器并且用filter(lambda p: p.requires_grad, model.parameters())显式过滤。3.2 参数分组不同层用不同学习率的落地写法分层学习率在微调场景几乎是标配。实现方式是给优化器传参数组backbone_params [] head_params [] for name, param in model.named_parameters(): if not param.requires_grad: continue if name.startswith(backbone): backbone_params.append(param) else: head_params.append(param) optimizer torch.optim.AdamW([ {params: backbone_params, lr: 1e-5}, {params: head_params, lr: 1e-3}, ], weight_decay0.01)这里要注意两点。第一weight_decay如果对不同组有不同要求要在每个组的字典里单独写顶层传的值会被组内覆盖。第二参数分组时务必确认没有参数被重复加入多个组否则优化器行为未定义。用named_parameters()遍历一次、按条件分桶是最安全的写法不要用两次独立的model.backbone.parameters()和model.head.parameters()去拼万一有共享参数就会重复。3.3 .to(device) 之后参数去哪了以及为什么不能缓存张量引用.to(device)是原地修改模块内部张量的设备不是返回一个新模型虽然它返回self。这意味着如果你在调用.to()之前把某个参数的引用存到了别处那个引用指向的还是旧设备上的张量。看这段model Model() w_ref model.fc1.weight model.to(cuda) print(w_ref.device) # cpu print(model.fc1.weight.device) # cudaw_ref还是 CPU 上的。所以不要在模型迁移设备之前缓存参数或 buffer 的引用需要时现取。这个坑在写自定义优化逻辑或手动更新参数时特别容易踩比如你提前存了self.weight到某个外部变量迁移后手动更新就更新到了旧张量上模型看起来“不学习”。buffer 同理。前面提过用register_buffer注册的常量会跟着.to()走但如果你在forward里把它取出来存成局部变量跨批次复用同样会遇到设备问题。原则很简单每次forward现取现用不跨调用缓存张量。4. 训练态与评估态train/eval 到底改了什么4.1 train() 和 eval() 是递归开关不是“训练/推理”的语义标签很多人以为model.train()是“开始训练”的指令model.eval()是“开始推理”的指令。严格说它们只是递归地把所有子模块的self.training标志设为 True 或 False具体行为由各层自己决定。Dropout 在trainingTrue时随机置零BatchNorm 在trainingTrue时用当前 batch 统计量并更新 running stats在False时用 running stats。仅此而已。理解这一点后很多“为什么 eval 结果和 train 结果差很多”的疑惑就有了解释不是模型坏了是 Dropout 和 BN 的行为本来就不同。评估时必须model.eval()并且通常配合torch.no_grad()减少显存占用。两者作用不同eval()改行为no_grad()关梯度记录评估时一般都要。4.2 自定义层如何正确响应 training 标志如果你写自定义层且行为依赖训练/评估态要在forward里读self.trainingclass StochasticDepth(nn.Module): def __init__(self, p0.1): super().__init__() self.p p def forward(self, x): if not self.training or self.p 0.0: return x keep 1 - self.p mask torch.empty(x.size(0), 1, 1, 1, devicex.device).bernoulli_(keep) return x * mask / keep关键是self.training会被父模块的train()/eval()递归设置你不需要手动维护。但要注意如果你在forward里用了torch.rand之类的随机操作却没判断self.training那评估时也会有随机性结果不可复现。这是自定义层最常见的疏漏之一。4.3 部分模块单独切换态的场景有些复杂模型需要“整体 eval 但某个子模块保持 train”比如某些对比学习框架里动量编码器需要特殊处理。做法是单独对子模块调用model.eval() model.momentum_encoder.train()这种写法要非常小心因为它打破了“整体一致”的假设。用之前想清楚那个子模块里的 BN 是否真的需要更新统计量Dropout 是否真的需要保留大多数情况下这种需求源于对 BN 行为的误判而不是真的需要混合态。我的建议是除非有明确的算法依据否则不要混用。5. state_dict 与模型保存那些保存了却加载不回来的坑5.1 state_dict 里到底有什么为什么它比 pickle 整个模型更可靠state_dict()返回一个从参数名到张量的字典包含所有Parameter和buffer但不包含模型结构代码。保存它、加载时先重建结构再load_state_dict是官方推荐做法。相比之下torch.save(model)直接 pickle 整个对象依赖类定义路径换目录、改类名、重构代码后极易加载失败。我在实际项目里一律只存state_dict配合一份结构配置可移植性好得多。参数名的层级结构值得留意。state_dict的键是带点号的路径比如backbone.blocks.0.proj.weight。这个路径由属性名和ModuleList的索引拼成。所以给子模块起名要稳定一旦改了属性名旧权重就加载不上。重构模型时如果必须改名要么写映射脚本转换键名要么用strictFalse加载后手动核对缺失和多余的键。5.2 load_state_dict 的 strict 参数与键不匹配排查load_state_dict默认strictTrue键集合必须完全一致多一个少一个都报错。排查键不匹配的标准流程是missing, unexpected model.load_state_dict(ckpt, strictFalse) print(missing:, missing) print(unexpected:, unexpected)missing是模型有但权重里没有的键unexpected是权重里有但模型没有的键。常见原因有三类一是模型结构改了加了层或删了层二是保存时用了DataParallel或分布式包装键名多了module.前缀三是自定义层里用了普通属性存张量没进state_dict导致保存时缺失。第二类特别常见。解决办法是加载前对键做前缀剥离new_ckpt {k.replace(module., , 1): v for k, v in ckpt.items()} model.load_state_dict(new_ckpt)反过来如果要把单卡权重加载进包装过的模型就手动加前缀。这类转换建议写成小工具函数别每次手敲。5.3 优化器状态也要存断点续训别只存模型只存模型权重、不存优化器状态续训时动量、二阶矩全部重置loss 曲线会出现明显跳变。完整断点应该包含模型state_dict、优化器state_dict、当前 epoch、学习率调度器状态以及随机数状态如果要求严格可复现torch.save({ model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), epoch: epoch, }, ckpt.pt)加载时对应地逐个load_state_dict。注意优化器的state_dict里保存的是参数索引到状态的映射它依赖参数组的顺序。所以加载优化器状态前必须先按保存时相同的顺序构建参数组否则状态会错位。这也是为什么参数分组逻辑要写成可复用的函数训练脚本和恢复脚本共用同一份避免顺序不一致。6. 几个高频“玄学问题”的定位思路6.1 参数没更新从登记、梯度、优化器三步查遇到“训练半天参数不变”按这个顺序排查基本能覆盖九成情况。第一步查登记len(list(model.parameters()))是否符合预期如果明显偏少多半是子模块没被正确登记回去检查是不是用了普通 list/dict。第二步查梯度对一个 batch 反传后打印param.grad是否为None或全零。为None说明该参数没参与计算图可能是forward里没用它或者被detach了全零可能是初始化或激活函数导致。第三步查优化器确认优化器持有的参数和模型参数是同一批对象冻结逻辑没有误伤学习率不是零。这三步我做成过一个检查清单每次新模型跑不通就照着过一遍比盲目改超参高效得多。6.2 显存越跑越大buffer 和计算图残留显存缓慢增长通常有两个来源。一是把带梯度的张量存进了列表跨批次累积计算图一直没释放。典型错误是在训练循环里losses.append(loss)而不是losses.append(loss.item())后者只存标量前者把整个计算图都留住了。二是自定义 buffer 在forward里被反复重新赋值且尺寸增长比如错误地self.cache torch.cat([self.cache, x])。buffer 应该尺寸固定需要累积统计量时用滑动平均而不是无限拼接。排查手段是定期打印torch.cuda.memory_allocated()配合在可疑位置插del加torch.cuda.empty_cache()做二分定位。empty_cache只是释放缓存分配器里没用的块不解决计算图残留别指望它治本。6.3 多卡包装后保存加载对不上用DataParallel或分布式包装后模型被套了一层state_dict键多了前缀。保存时如果存的是包装后的模型键就带前缀加载时如果重建的是未包装模型就对不上。统一策略是保存时存model.module.state_dict()如果是包装过的加载时先建未包装模型再加载。这样权重文件与包装方式解耦换卡数、换包装策略都不用改权重。分布式场景下还要注意只有 rank 0 应该执行保存其他 rank 跳过否则多进程同时写同一文件会损坏。加载时所有 rank 都要加载保证初始权重一致。7. 写 Module 时我逐渐固化下来的几条习惯第一条任何装子模块的容器只用nn提供的三种普通 list/dict 一律不碰。这条帮我省掉了大量排查时间。第二条需要跨设备、跨保存的张量一律register_buffer普通属性只放配置值int、float、字符串、tuple。配置值不进state_dict是合理的因为它们由代码决定但凡是张量就该被管理。第三条forward保持纯粹只做张量运算不建层、不改结构、不做 I/O。需要条件分支的结构差异在__init__里用配置决定好forward里只做数据路由。第四条保存只存state_dict加配置不 pickle 整个模型。断点续训把优化器和调度器状态一起存参数分组逻辑抽成函数复用。第五条新模型先跑一遍参数登记和梯度检查确认parameters()数量、名字符合预期再开始正式训练。这个前置检查花不了两分钟但能避免几小时的无效训练。这些习惯谈不上多高深但都是从实际返工里换来的。Module这个类本身不复杂复杂的是它背后那套“登记—管理—序列化”的约定把这些约定吃透写模型时心里就有底出了问题也知道往哪查。