我见过不少从别的框架转过来的朋友第一次在PyTorch里写模型都会被nn.Module搞懵。明明是几个nn.Linear叠在一起为什么必须写一个类为什么forward里不能随便建层为什么模型保存后再加载有时候参数能对上有时候又报错说缺了某个key这些问题看起来零散但根子都在同一个地方没有真正理解nn.Module是怎么管理模型参数的。这篇内容就围绕模型创建和nn.Module展开目标是把一个模型对象到底是怎么被组装出来的讲透。它适合刚入门PyTorch的读者也适合已经写了一阵子模型、但遇到诡异问题只能靠搜索引擎救火的同学。读完你会明白nn.Module不只是一个装网络层的容器它更像一个登记系统管着层的挂载、参数的收集、设备的迁移、训练和推理状态的切换还有钩子回调。理解了这套机制你后面调模型、改模型、自定义层都会顺手很多。1. nn.Module的底层设计它凭什么连接起网络层和可训练参数1.1 每一个普通赋值背后都有一本登记簿nn.Module的底层不复杂但设计非常巧妙。它的核心手段是重写了Python的__setattr__方法也就是所有self.xxx yyy这种赋值操作都会先经过它检查。当你写出self.fc nn.Linear(64, 128)这一行代码时nn.Module在背后做的事情是先检查被赋值的对象是不是nn.Module类型如果是就把这个对象放进当前模块的_modules字典里。_modules字典的key就是属性名fcvalue就是这个Linear实例。同样的机制也作用于参数。你写self.weight nn.Parameter(torch.randn(64, 128))nn.Module会识别出nn.Parameter类型把对象放进_parameters字典。这里有个非常重要的隐含条件只有直接赋值给__init__里的self.属性名nn.Module才能捕获到。这也解释了一个高频报错——忘写super().__init__()。我见过很多次class MyModel(nn.Module): def __init__(self): # 忘记了 super().__init__() self.fc nn.Linear(32, 64)当时看起来一切正常模型能创建forward也能跑但一调用model.parameters()就会发现参数列表是空的。原因很简单没有调用父类构造方法_modules和_parameters这些字典根本不存在赋值操作自然落不到登记簿上。而且这类错误在训练开始前往往暴露不出来直到优化器里传了空参数列表loss完全不下降才让人意识到出了问题。所以我的建议是写任何继承nn.Module的类第一时间先写super().__init__()把这个登记簿初始化好再去挂层、挂参数。这个顺序习惯养成之后能帮你避开一大类莫名其妙的bug。1.2 用list包层的致命错误为什么必须使用ModuleList/ModuleDict明白了上面的登记机制一个非常经典的坑就很容易理解了。很多人写模型时图方便会这样写class BadMLP(nn.Module): def __init__(self): super().__init__() self.layers [nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10)] def forward(self, x): for layer in self.layers: x layer(x) return x model BadMLP() print(model.state_dict().keys()) # odict_keys([])forward能正常跑但模型里一个可训练参数都没有state_dict()是空的。原因就出在登记机制上self.layers被赋值的是一个Python原生listnn.Module只会把子模块放入_modules字典对于普通list它会认为这就是一个普通属性根本不去检查里面的元素。更隐蔽的是这个list里的Linear层虽然持有自己的参数但它们和父模块之间没有建立任何登记关系父模块的parameters()方法自然遍历不到它们。优化器不会更新它们保存模型时也会丢失它们。运行时不报错结果却是模型完全没学到东西。正确的写法是换成nn.ModuleListself.layers nn.ModuleList([ nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10) ])nn.ModuleList是专门为存放多个子模块准备的容器类它同样重写了赋值和迭代逻辑核心保证是无论list里装了多少个层这些层都会逐一登记到父模块的_modules字典里父模块的parameters()才能收集到它们。同样的规则也适用于nn.ModuleDict。当你需要动态管理一批模块用字典按名字组织时应该用nn.ModuleDict而不是原生dict。我在做多分支结构时用过它比如给不同输入来源分配编码器用nn.ModuleDict后可以非常方便地通过字符串key动态取用模块而且参数注册完全不会丢。还有一个容易被忽略的场景如果你确实需要在__init__里用Python list也可以手动往_modules里塞。但不推荐代码会变得很绕可读性差。正确的工具就是nn.ModuleList和nn.ModuleDict它们的存在就是为了解决这个问题。2. 模型创建的正确姿势__init__做装配forward定义计算流程2.1 Sequential适合一条直线的网络但条件必须清楚每次看新手写的模型最常见的形式就是model nn.Sequential( nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10) )nn.Sequential用起来确实爽把层按照执行顺序排好forward的时候自动一个个传下去。内部实现其实也不神秘Sequential自己也是一个nn.Module它把传入的每个层按顺序放进_modules字典key是递增数字0,1,2...forward时遍历字典按顺序调用。所以它的参数收集、设备迁移、save/load这些行为和手写子模块的模型完全一致。这个写法适合网络结构是一条直线的模型输入经过若干层输出中途没有任何分支和跳连。第一版模型用Sequential快速验证想法完全没有问题。我个人的习惯是当模型结构简单、后处理不多的时候直接Sequential一旦出现多层共享输入不同分支最后汇合需要根据条件动态选择不同路径这类需求就不再硬套Sequential了因为硬套会让代码变成把所有逻辑塞进一个函数反而更难看。2.2 手写forward的场景多输入、分支结构、动态深度nn.Module真正灵活的地方在于它不限制forward里的写法。你继承nn.Module后forward方法就是一段普通的Python代码层可以是任意的计算步骤。举个多输入分支汇合的例子class TwoBranchModel(nn.Module): def __init__(self): super().__init__() self.text_encoder nn.Sequential(nn.Linear(32, 64), nn.ReLU()) self.image_encoder nn.Sequential(nn.Linear(64, 64), nn.ReLU()) self.fusion nn.Linear(64, 10) def forward(self, text_feat, image_feat): t self.text_encoder(text_feat) v self.image_encoder(image_feat) return self.fusion(t v)这里forward接受两个输入分别走不同编码器最后融合。类似的场景还有根据输入的某个标志位选择不同分支比如def forward(self, x, use_shortcutTrue): if use_shortcut: return self.fc2(self.fc1(x)) x return self.fc2(self.fc1(x))对于这类动态流程手写forward几乎是唯一的选择因为你没法用Sequential表达条件分支。还有一类场景是动态深度比如循环层数可配置class IterativeRefine(nn.Module): def __init__(self, hidden_dim, num_iters): super().__init__() self.num_iters num_iters self.layer nn.Linear(hidden_dim, hidden_dim) def forward(self, x): for _ in range(self.num_iters): x torch.relu(self.layer(x)) return x所以我的习惯是第一版模型默认手写类把层放在__init__里forward做最小化的事。这样后续加模块、加分支、加辅助损失都方便。Sequential更多用于子模块内部的快速堆叠整体结构用类来组织两者并不冲突。2.3 新手最容易犯的错在forward里new一个新层这是我在代码评审里经常抓住的问题。有人图省事直接在forward里创建层class BadPattern(nn.Module): def __init__(self): super().__init__() def forward(self, x): layer nn.Linear(x.size(-1), 10) return layer(x)表面上看起来完全能跑。但每跑一次forward都会new一个全新的Linear里面的权重是随机初始化的。这就导致两个问题第一这个层从来没有参与过任何优化步骤它的权重始终是随机值模型输出没有任何意义第二更隐蔽的是每次forward还会往模型参数里添加新的Parameter多次前向之后模型保存的参数量会异常膨胀还可能出现加载模型时key对不上的诡异现象。正确的做法永远是把层放进__init__里forward只负责调用。forward是计算过程不是装配过程。这个原则不仅适用于nn.Module也适用于任何组件化的代码设计创建对象一次调用多次。3. 藏在文档字缝里的机制参数、状态字典与调试钩子3.1 parameters、named_parameters和state_dict到底有什么区别模型创建完之后我们需要把参数交给优化器或者保存到文件或者加载预训练权重。对应地nn.Module提供了几个看起来很相似的方法parameters()、named_parameters()、state_dict()。它们在底层查的是同一批登记数据但返回的信息粒度不一样。我自己整理过一张对比方法返回内容典型用途model.parameters()按注册顺序产出所有叶子参数传给优化器如optim.SGD(model.parameters())model.named_parameters()产出一对(name, param)name是带路径的字符串查看结构比如fc1.weight、fc2.biasmodel.state_dict()返回一个字典包含参数和bufferkey与named_parameters对应保存、加载、初始化、查看全部权重named_parameters()里的name格式值得注意。它体现的是模块树结构。比如一个模型里有个encoder子模块它的fc层的权重name就是encoder.fc.weight点号路径是从顶层一路往下的属性名拼接。这个格式不仅方便打印查看也直接对应state_dict()里的key。加载权重时出现size mismatch或者missing key的报错就是拿这个key和预训练文件的key做匹配时出了问题。parameters()方法遍历的是_parameters里面所有叶子参数包括嵌套在子模块里的。它有recurseTrue参数默认递归收集整棵模块树。如果你只想拿某个子模块的参数也可以直接访问子模块的parameters()方法。3.2 to(device)、train()/eval()如何递归传遍整棵模型树刚接触PyTorch的时候我一度以为model.to(device)就是遍历参数挨个调用.to()。后来发现它的机制是递归的nn.Module调用_apply方法对所有注册的Parameter、buffer以及所有子模块统一调用.to()保证整个模块树里的张量都迁移到同一个设备上。这里有一个很容易忽略的行为model.to(device)是inplace地修改了模型内部张量的内存位置但它同时也返回模型自身的引用。所以写model model.to(device)和直接model.to(device)效果一样。很多官方代码习惯写成model.to(device)不加返回值因为模型内部的参数已经被迁移了。model.train()和model.eval()的传播逻辑和to(device)完全一致都是递归地遍历子模块切换每个子模块的训练状态。它们内部维护的是每个模块的training布尔标志。这个标志主要影响两类层Dropout是否启用、BatchNorm是否使用当前batch统计量来归一化。我见过很多新手以为model.eval()是关闭模型计算图或者停止梯度传播其实完全不是。eval()模式下模型前向仍然参与autograd只是部分层的行为变了。如果你想真正关闭梯度计算要在torch.no_grad()上下文里单独做。3.3 forward hook不改源码就能偷看每一层的输出模型调试最直接的需求是我想看看中间层输出到底是什么形状、什么数值分布。最朴素的办法是在forward里加print但这要改动源码加进去还得删掉非常麻烦。nn.Module的hook机制就是为此设计的。典型的用法是def hook_fn(module, input, output): print(flayer name: {module} | input shape: {input[0].shape} | output shape: {output.shape}) handle model.fc1.register_forward_hook(hook_fn) # 跑完forward后 handle.remove()register_forward_hook返回一个句柄想移除hook时调用handle.remove()即可。它接收的回调函数有三个参数当前模块、输入张量元组、输出张量。你可以在这个函数里做任意检查比如统计输出是否含有NaN、打印分布式状态、记录激活值的均值方差。这套机制在跑复杂大模型时非常实用。我调试过一个多层叠加的模型某一层输出爆炸连带着后面全部变成NaN。用hook在所有层上注册同一个小函数打印每一层的输出统计量就能精确锁定是第几层出的问题。这是我在看中间层输出这个需求上最常用也最推荐的方法不需要改任何业务代码。4. 我调试nn.Module时踩过的典型坑完整排查链路4.1 模型丢层state_dict里突然少了一截key某次项目里我给模型加了一个辅助分类头结构上明明写得好好的forward也调用了但保存模型时发现state_dict()里根本没有这个分类头的参数。排查链路是这样的第一步先打印model看结构确认模块显示正常。第二步打印model.state_dict().keys()发现确实少了aux_head相关的key。第三步回到__init__检查赋值语句。最后发现我把辅助头放在了一个普通Python list里self.aux_heads [nn.Linear(128, 10), nn.Linear(128, 5)]问题和我前面讲的一模一样。这种情况在__init__结构变复杂后尤其容易发生层多了以后用list做分组是人的本能但nn.Module不认识原生list。排查思路总结起来就是缺参数先看结构结构问题再看赋值方式赋值方式最后确认容器类型。只要遇到运行正常但参数缺失的问题优先怀疑容器是否合法。4.2 所有张量必须同设备device不一致问题的定位路径训练刚开始报了一个经典错误RuntimeError: Expected all tensors to be on the same device。看堆栈定位到forward里某个加法操作分别打印两个操作数的设备发现一个是cuda:0一个是cpu。深挖下去原因是模型里有一个子模块是在forward里用x.device临时构造的而这个临时构造的参数虽然赋给了self但它是在forward首次运行时才注册进去的。这类问题的处理链路我一般按三步走在训练脚本最开始把模型model.to(device)先执行。检查输入数据batch batch.to(device)这一步在Dataloader返回后立刻做。遍历模型的named_parameters()和named_buffers()打印各自的device找出漏网之鱼。真正治本的做法还是我之前强调的所有需要训练的层和参数在__init__阶段就创建好forward只做计算。这样to(device)只需要在初始阶段执行一次整棵模型树上的张量位置就都一致了。4.3 反向传播时梯度为None原地操作留下的暗伤某个模型跑起来loss很正常但backward之后所有参数梯度都是None优化器一更新就报错。查了半天发现我在一个自定义层里图省事用了原地操作修改权重self.weight / factorweight / factor本质上是weight.__itruediv__(factor)这种原地修改会破坏autograd的计算图记录。PyTorch的自动微分依赖计算图记录每次操作的依赖关系原地修改相当于把图中的一个节点直接覆盖掉了梯度回溯时找不到原来的依赖路径自然就无法计算梯度。排查这类问题的标准做法是先开启异常检测torch.autograd.set_detect_anomaly(True)再跑一次训练让PyTorch给出更精确的定位信息同时配合在每个可疑层注册backward hook打印grad是否为None。定位到具体层之后把原地操作改成赋值式self.weight self.weight / factor或者其他不破坏计算图的方式。这里有个原则对参与梯度计算的parameter尽量不要用inplace操作除非你能明确知道自己在做什么。4.4 共享参数的惊喜同一个层对象被多处复用有次我写一个双塔结构的模型想把同一个编码器放在两条塔里共享。我最初写的是self.shared nn.Linear(64, 64) self.left_head self.shared self.right_head self.shared运行起来完全正常说明我的账面上得到了一个共享模型。但我测了一下发现left_head和right_head的两个层地址完全相同它们本来就是同一个对象。PyTorch的参数管理是基于对象引用的不是复制。self.left_head self.shared只是给同一个对象起了一个新名字并没有复制权重。很多真正需要共享的地方这个行为是好用的比如孪生网络的正规写法。但有时候你本意是想复制一个结构相同、权重独立的层却意外写成共享就会出现奇怪的结果修改一个塔的权重另一个也跟着变。排查链路是打印id(model.left_head),id(model.right_head)如果相同就是同一对象。如果确实需要深度复制一个结构相同但参数独立的模块正确做法是先copy.deepcopy(model.shared)再重新包装成一个子模块。注意不能在__init__里直接self.left_head copy.deepcopy(self.shared)之前先给self.shared赋值否则赋值者又会触发登记问题。流程是先创建原始层、复制、再挂着三个名字并且每个都必须直接赋值给self属性。4.5 多卡训练时不要手动把batch拆给模型还有一个高频错误用单卡逻辑编写代码然后手动把batch切一半分别塞给model再想把梯度拼起来。这个操作在多卡时代基本是白费功夫因为DataParallel和DistributedDataParallel已经帮你做了batch拆分、梯度同步这些事情。自己手动拆很容易出现部分张量在不同device的混乱问题还会让代码里到处都是device判断。正确做法是始终把整个batch放主卡交给DDP封装后的模型去forward框架会自动把输入按维度分片到各卡再汇总输出。如果你写的模型里含buffer比如BatchNorm的running统计量多卡注意同步buffer选项通常把broadcast_buffers设为True即可。5. 进阶实践把自定义模块写得像官方模块一样可靠5.1 自定义层的完整骨架Parameter与buffer的注册规则自己动手写一个层是检验nn.Module理解程度的最好方式。拿一个简单的可学习加权层举例class WeightedSumLayer(nn.Module): def __init__(self, dim): super().__init__() self.gamma nn.Parameter(torch.ones(dim)) self.register_buffer(running_mean, torch.zeros(dim)) def forward(self, x): self.running_mean.copy_(x.mean(dim0)) return x * self.gamma self.running_mean这里面有两条规则需要区分清楚nn.Parameter包装的是想被优化器更新的权重会出现在parameters()里。register_buffer注册的是不参与优化、但需要随模型保存和迁移的状态比如BatchNorm的running_mean、running_var以及一些EMA统计量。它不会出现在parameters()里但会出现在state_dict()的buffer部分。很多人会在自定义层里直接写self.some_tensor torch.randn(...)结果在state_dict()里看到一个奇怪的key还出现在buffers中。这是nn.Module的另一个默认行为普通tensor赋值给self.xxx会被自动当作buffer注册。如果你确实想保存这个状态这没问题但如果你只是想在forward里用一下临时tensor就不该挂到self上面应该先创建后传入或者用torch.zeros这种普通局部变量否则state_dict会越来越臃肿。5.2 用hook定位哪一层输出变成了NaN模型训练到一半loss突然变成NaN是每个人都遇到过的经典问题。我有一套固定打法先用hook在所有层上注册检查函数一旦发现输出中有NaN立刻打印层的名字、输入输出统计量并终止。def nan_check_fn(module, input, output): if torch.isnan(output).any(): print(fNaN detected in {module}) raise SystemExit handles [] for name, submodule in model.named_modules(): if isinstance(submodule, nn.Module) and hasattr(submodule, forward): handles.append(submodule.register_forward_hook(nan_check_fn))跑一次后第一个触发NaN的就是问题层。接下来再检查它的输入看是上一层的输出问题还是这一层自身的权重问题。这种排查方式不需要改forward里的任何代码通常几分钟就能定位到具体位置。需要注意的是named_modules()会递归遍历整棵模型树对每个子模块注册hook可能让前向变慢。但作为定位手段这点开销完全值得。定位完成后记得把handles全部remove()否则后续训练一直开着检查。5.3 让模型代码持续可维护的几个习惯最后分享几个我长期写PyTorch模型的习惯这些习惯都是从踩坑里提炼出来的第一个习惯所有层都在__init__里装配完成forward里只做逻辑编排。这样model.parameters()、state_dict()、to(device)、train()/eval()这些机制都不需要额外照顾所有状态自动一致。第二个习惯统一参数初始化。在__init__最后写一个self.apply(self._init_weights)把初始化逻辑集中在专门的函数里利用apply递归调用到所有子模块上的特性来一遍初始化def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.kaiming_uniform_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias)第三个习惯写完一个新模型先打印三样东西模型结构、state_dict()的key列表、以及喂一个随机tensor会不会正常forward和backward。前两步是验证装好了第三步是验证能跑通。我靠这个简单流程在把模型交付给训练脚本之前几乎可以过滤掉90%的常见问题。第四个习惯谨慎使用Python原生的list和dict存模块。在模型定义这个场景里原生容器只有在不涉及参数管理时才安全只要容器里的东西是nn.Module或nn.Parameter一律用ModuleList、ModuleDict、nn.Sequential。说回nn.Module官方文档其实很短但它设计的精妙程度远超过表面上那几百行代码。把赋值、收集、递归这些基础机制弄清楚之后你会发现模型创建这件事没有太多魔法可言——它就是一套严格的登记和递归规则。自己动手写一个自定义层再试着用hook查一次NaN出现的位置这套实战走下来比背十遍文档印象都深。 
长河君 · 2026/10/15 17:07:43 · 阅读 1.2k 深入理解PyTorch nn.Module：参数登记与模型创建的底层机制
我见过不少从别的框架转过来的朋友第一次在PyTorch里写模型都会被nn.Module搞懵。明明是几个nn.Linear叠在一起为什么必须写一个类为什么forward里不能随便建层为什么模型保存后再加载有时候参数能对上有时候又报错说缺了某个key这些问题看起来零散但根子都在同一个地方没有真正理解nn.Module是怎么管理模型参数的。这篇内容就围绕模型创建和nn.Module展开目标是把一个模型对象到底是怎么被组装出来的讲透。它适合刚入门PyTorch的读者也适合已经写了一阵子模型、但遇到诡异问题只能靠搜索引擎救火的同学。读完你会明白nn.Module不只是一个装网络层的容器它更像一个登记系统管着层的挂载、参数的收集、设备的迁移、训练和推理状态的切换还有钩子回调。理解了这套机制你后面调模型、改模型、自定义层都会顺手很多。1. nn.Module的底层设计它凭什么连接起网络层和可训练参数1.1 每一个普通赋值背后都有一本登记簿nn.Module的底层不复杂但设计非常巧妙。它的核心手段是重写了Python的__setattr__方法也就是所有self.xxx yyy这种赋值操作都会先经过它检查。当你写出self.fc nn.Linear(64, 128)这一行代码时nn.Module在背后做的事情是先检查被赋值的对象是不是nn.Module类型如果是就把这个对象放进当前模块的_modules字典里。_modules字典的key就是属性名fcvalue就是这个Linear实例。同样的机制也作用于参数。你写self.weight nn.Parameter(torch.randn(64, 128))nn.Module会识别出nn.Parameter类型把对象放进_parameters字典。这里有个非常重要的隐含条件只有直接赋值给__init__里的self.属性名nn.Module才能捕获到。这也解释了一个高频报错——忘写super().__init__()。我见过很多次class MyModel(nn.Module): def __init__(self): # 忘记了 super().__init__() self.fc nn.Linear(32, 64)当时看起来一切正常模型能创建forward也能跑但一调用model.parameters()就会发现参数列表是空的。原因很简单没有调用父类构造方法_modules和_parameters这些字典根本不存在赋值操作自然落不到登记簿上。而且这类错误在训练开始前往往暴露不出来直到优化器里传了空参数列表loss完全不下降才让人意识到出了问题。所以我的建议是写任何继承nn.Module的类第一时间先写super().__init__()把这个登记簿初始化好再去挂层、挂参数。这个顺序习惯养成之后能帮你避开一大类莫名其妙的bug。1.2 用list包层的致命错误为什么必须使用ModuleList/ModuleDict明白了上面的登记机制一个非常经典的坑就很容易理解了。很多人写模型时图方便会这样写class BadMLP(nn.Module): def __init__(self): super().__init__() self.layers [nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10)] def forward(self, x): for layer in self.layers: x layer(x) return x model BadMLP() print(model.state_dict().keys()) # odict_keys([])forward能正常跑但模型里一个可训练参数都没有state_dict()是空的。原因就出在登记机制上self.layers被赋值的是一个Python原生listnn.Module只会把子模块放入_modules字典对于普通list它会认为这就是一个普通属性根本不去检查里面的元素。更隐蔽的是这个list里的Linear层虽然持有自己的参数但它们和父模块之间没有建立任何登记关系父模块的parameters()方法自然遍历不到它们。优化器不会更新它们保存模型时也会丢失它们。运行时不报错结果却是模型完全没学到东西。正确的写法是换成nn.ModuleListself.layers nn.ModuleList([ nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10) ])nn.ModuleList是专门为存放多个子模块准备的容器类它同样重写了赋值和迭代逻辑核心保证是无论list里装了多少个层这些层都会逐一登记到父模块的_modules字典里父模块的parameters()才能收集到它们。同样的规则也适用于nn.ModuleDict。当你需要动态管理一批模块用字典按名字组织时应该用nn.ModuleDict而不是原生dict。我在做多分支结构时用过它比如给不同输入来源分配编码器用nn.ModuleDict后可以非常方便地通过字符串key动态取用模块而且参数注册完全不会丢。还有一个容易被忽略的场景如果你确实需要在__init__里用Python list也可以手动往_modules里塞。但不推荐代码会变得很绕可读性差。正确的工具就是nn.ModuleList和nn.ModuleDict它们的存在就是为了解决这个问题。2. 模型创建的正确姿势__init__做装配forward定义计算流程2.1 Sequential适合一条直线的网络但条件必须清楚每次看新手写的模型最常见的形式就是model nn.Sequential( nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10) )nn.Sequential用起来确实爽把层按照执行顺序排好forward的时候自动一个个传下去。内部实现其实也不神秘Sequential自己也是一个nn.Module它把传入的每个层按顺序放进_modules字典key是递增数字0,1,2...forward时遍历字典按顺序调用。所以它的参数收集、设备迁移、save/load这些行为和手写子模块的模型完全一致。这个写法适合网络结构是一条直线的模型输入经过若干层输出中途没有任何分支和跳连。第一版模型用Sequential快速验证想法完全没有问题。我个人的习惯是当模型结构简单、后处理不多的时候直接Sequential一旦出现多层共享输入不同分支最后汇合需要根据条件动态选择不同路径这类需求就不再硬套Sequential了因为硬套会让代码变成把所有逻辑塞进一个函数反而更难看。2.2 手写forward的场景多输入、分支结构、动态深度nn.Module真正灵活的地方在于它不限制forward里的写法。你继承nn.Module后forward方法就是一段普通的Python代码层可以是任意的计算步骤。举个多输入分支汇合的例子class TwoBranchModel(nn.Module): def __init__(self): super().__init__() self.text_encoder nn.Sequential(nn.Linear(32, 64), nn.ReLU()) self.image_encoder nn.Sequential(nn.Linear(64, 64), nn.ReLU()) self.fusion nn.Linear(64, 10) def forward(self, text_feat, image_feat): t self.text_encoder(text_feat) v self.image_encoder(image_feat) return self.fusion(t v)这里forward接受两个输入分别走不同编码器最后融合。类似的场景还有根据输入的某个标志位选择不同分支比如def forward(self, x, use_shortcutTrue): if use_shortcut: return self.fc2(self.fc1(x)) x return self.fc2(self.fc1(x))对于这类动态流程手写forward几乎是唯一的选择因为你没法用Sequential表达条件分支。还有一类场景是动态深度比如循环层数可配置class IterativeRefine(nn.Module): def __init__(self, hidden_dim, num_iters): super().__init__() self.num_iters num_iters self.layer nn.Linear(hidden_dim, hidden_dim) def forward(self, x): for _ in range(self.num_iters): x torch.relu(self.layer(x)) return x所以我的习惯是第一版模型默认手写类把层放在__init__里forward做最小化的事。这样后续加模块、加分支、加辅助损失都方便。Sequential更多用于子模块内部的快速堆叠整体结构用类来组织两者并不冲突。2.3 新手最容易犯的错在forward里new一个新层这是我在代码评审里经常抓住的问题。有人图省事直接在forward里创建层class BadPattern(nn.Module): def __init__(self): super().__init__() def forward(self, x): layer nn.Linear(x.size(-1), 10) return layer(x)表面上看起来完全能跑。但每跑一次forward都会new一个全新的Linear里面的权重是随机初始化的。这就导致两个问题第一这个层从来没有参与过任何优化步骤它的权重始终是随机值模型输出没有任何意义第二更隐蔽的是每次forward还会往模型参数里添加新的Parameter多次前向之后模型保存的参数量会异常膨胀还可能出现加载模型时key对不上的诡异现象。正确的做法永远是把层放进__init__里forward只负责调用。forward是计算过程不是装配过程。这个原则不仅适用于nn.Module也适用于任何组件化的代码设计创建对象一次调用多次。3. 藏在文档字缝里的机制参数、状态字典与调试钩子3.1 parameters、named_parameters和state_dict到底有什么区别模型创建完之后我们需要把参数交给优化器或者保存到文件或者加载预训练权重。对应地nn.Module提供了几个看起来很相似的方法parameters()、named_parameters()、state_dict()。它们在底层查的是同一批登记数据但返回的信息粒度不一样。我自己整理过一张对比方法返回内容典型用途model.parameters()按注册顺序产出所有叶子参数传给优化器如optim.SGD(model.parameters())model.named_parameters()产出一对(name, param)name是带路径的字符串查看结构比如fc1.weight、fc2.biasmodel.state_dict()返回一个字典包含参数和bufferkey与named_parameters对应保存、加载、初始化、查看全部权重named_parameters()里的name格式值得注意。它体现的是模块树结构。比如一个模型里有个encoder子模块它的fc层的权重name就是encoder.fc.weight点号路径是从顶层一路往下的属性名拼接。这个格式不仅方便打印查看也直接对应state_dict()里的key。加载权重时出现size mismatch或者missing key的报错就是拿这个key和预训练文件的key做匹配时出了问题。parameters()方法遍历的是_parameters里面所有叶子参数包括嵌套在子模块里的。它有recurseTrue参数默认递归收集整棵模块树。如果你只想拿某个子模块的参数也可以直接访问子模块的parameters()方法。3.2 to(device)、train()/eval()如何递归传遍整棵模型树刚接触PyTorch的时候我一度以为model.to(device)就是遍历参数挨个调用.to()。后来发现它的机制是递归的nn.Module调用_apply方法对所有注册的Parameter、buffer以及所有子模块统一调用.to()保证整个模块树里的张量都迁移到同一个设备上。这里有一个很容易忽略的行为model.to(device)是inplace地修改了模型内部张量的内存位置但它同时也返回模型自身的引用。所以写model model.to(device)和直接model.to(device)效果一样。很多官方代码习惯写成model.to(device)不加返回值因为模型内部的参数已经被迁移了。model.train()和model.eval()的传播逻辑和to(device)完全一致都是递归地遍历子模块切换每个子模块的训练状态。它们内部维护的是每个模块的training布尔标志。这个标志主要影响两类层Dropout是否启用、BatchNorm是否使用当前batch统计量来归一化。我见过很多新手以为model.eval()是关闭模型计算图或者停止梯度传播其实完全不是。eval()模式下模型前向仍然参与autograd只是部分层的行为变了。如果你想真正关闭梯度计算要在torch.no_grad()上下文里单独做。3.3 forward hook不改源码就能偷看每一层的输出模型调试最直接的需求是我想看看中间层输出到底是什么形状、什么数值分布。最朴素的办法是在forward里加print但这要改动源码加进去还得删掉非常麻烦。nn.Module的hook机制就是为此设计的。典型的用法是def hook_fn(module, input, output): print(flayer name: {module} | input shape: {input[0].shape} | output shape: {output.shape}) handle model.fc1.register_forward_hook(hook_fn) # 跑完forward后 handle.remove()register_forward_hook返回一个句柄想移除hook时调用handle.remove()即可。它接收的回调函数有三个参数当前模块、输入张量元组、输出张量。你可以在这个函数里做任意检查比如统计输出是否含有NaN、打印分布式状态、记录激活值的均值方差。这套机制在跑复杂大模型时非常实用。我调试过一个多层叠加的模型某一层输出爆炸连带着后面全部变成NaN。用hook在所有层上注册同一个小函数打印每一层的输出统计量就能精确锁定是第几层出的问题。这是我在看中间层输出这个需求上最常用也最推荐的方法不需要改任何业务代码。4. 我调试nn.Module时踩过的典型坑完整排查链路4.1 模型丢层state_dict里突然少了一截key某次项目里我给模型加了一个辅助分类头结构上明明写得好好的forward也调用了但保存模型时发现state_dict()里根本没有这个分类头的参数。排查链路是这样的第一步先打印model看结构确认模块显示正常。第二步打印model.state_dict().keys()发现确实少了aux_head相关的key。第三步回到__init__检查赋值语句。最后发现我把辅助头放在了一个普通Python list里self.aux_heads [nn.Linear(128, 10), nn.Linear(128, 5)]问题和我前面讲的一模一样。这种情况在__init__结构变复杂后尤其容易发生层多了以后用list做分组是人的本能但nn.Module不认识原生list。排查思路总结起来就是缺参数先看结构结构问题再看赋值方式赋值方式最后确认容器类型。只要遇到运行正常但参数缺失的问题优先怀疑容器是否合法。4.2 所有张量必须同设备device不一致问题的定位路径训练刚开始报了一个经典错误RuntimeError: Expected all tensors to be on the same device。看堆栈定位到forward里某个加法操作分别打印两个操作数的设备发现一个是cuda:0一个是cpu。深挖下去原因是模型里有一个子模块是在forward里用x.device临时构造的而这个临时构造的参数虽然赋给了self但它是在forward首次运行时才注册进去的。这类问题的处理链路我一般按三步走在训练脚本最开始把模型model.to(device)先执行。检查输入数据batch batch.to(device)这一步在Dataloader返回后立刻做。遍历模型的named_parameters()和named_buffers()打印各自的device找出漏网之鱼。真正治本的做法还是我之前强调的所有需要训练的层和参数在__init__阶段就创建好forward只做计算。这样to(device)只需要在初始阶段执行一次整棵模型树上的张量位置就都一致了。4.3 反向传播时梯度为None原地操作留下的暗伤某个模型跑起来loss很正常但backward之后所有参数梯度都是None优化器一更新就报错。查了半天发现我在一个自定义层里图省事用了原地操作修改权重self.weight / factorweight / factor本质上是weight.__itruediv__(factor)这种原地修改会破坏autograd的计算图记录。PyTorch的自动微分依赖计算图记录每次操作的依赖关系原地修改相当于把图中的一个节点直接覆盖掉了梯度回溯时找不到原来的依赖路径自然就无法计算梯度。排查这类问题的标准做法是先开启异常检测torch.autograd.set_detect_anomaly(True)再跑一次训练让PyTorch给出更精确的定位信息同时配合在每个可疑层注册backward hook打印grad是否为None。定位到具体层之后把原地操作改成赋值式self.weight self.weight / factor或者其他不破坏计算图的方式。这里有个原则对参与梯度计算的parameter尽量不要用inplace操作除非你能明确知道自己在做什么。4.4 共享参数的惊喜同一个层对象被多处复用有次我写一个双塔结构的模型想把同一个编码器放在两条塔里共享。我最初写的是self.shared nn.Linear(64, 64) self.left_head self.shared self.right_head self.shared运行起来完全正常说明我的账面上得到了一个共享模型。但我测了一下发现left_head和right_head的两个层地址完全相同它们本来就是同一个对象。PyTorch的参数管理是基于对象引用的不是复制。self.left_head self.shared只是给同一个对象起了一个新名字并没有复制权重。很多真正需要共享的地方这个行为是好用的比如孪生网络的正规写法。但有时候你本意是想复制一个结构相同、权重独立的层却意外写成共享就会出现奇怪的结果修改一个塔的权重另一个也跟着变。排查链路是打印id(model.left_head),id(model.right_head)如果相同就是同一对象。如果确实需要深度复制一个结构相同但参数独立的模块正确做法是先copy.deepcopy(model.shared)再重新包装成一个子模块。注意不能在__init__里直接self.left_head copy.deepcopy(self.shared)之前先给self.shared赋值否则赋值者又会触发登记问题。流程是先创建原始层、复制、再挂着三个名字并且每个都必须直接赋值给self属性。4.5 多卡训练时不要手动把batch拆给模型还有一个高频错误用单卡逻辑编写代码然后手动把batch切一半分别塞给model再想把梯度拼起来。这个操作在多卡时代基本是白费功夫因为DataParallel和DistributedDataParallel已经帮你做了batch拆分、梯度同步这些事情。自己手动拆很容易出现部分张量在不同device的混乱问题还会让代码里到处都是device判断。正确做法是始终把整个batch放主卡交给DDP封装后的模型去forward框架会自动把输入按维度分片到各卡再汇总输出。如果你写的模型里含buffer比如BatchNorm的running统计量多卡注意同步buffer选项通常把broadcast_buffers设为True即可。5. 进阶实践把自定义模块写得像官方模块一样可靠5.1 自定义层的完整骨架Parameter与buffer的注册规则自己动手写一个层是检验nn.Module理解程度的最好方式。拿一个简单的可学习加权层举例class WeightedSumLayer(nn.Module): def __init__(self, dim): super().__init__() self.gamma nn.Parameter(torch.ones(dim)) self.register_buffer(running_mean, torch.zeros(dim)) def forward(self, x): self.running_mean.copy_(x.mean(dim0)) return x * self.gamma self.running_mean这里面有两条规则需要区分清楚nn.Parameter包装的是想被优化器更新的权重会出现在parameters()里。register_buffer注册的是不参与优化、但需要随模型保存和迁移的状态比如BatchNorm的running_mean、running_var以及一些EMA统计量。它不会出现在parameters()里但会出现在state_dict()的buffer部分。很多人会在自定义层里直接写self.some_tensor torch.randn(...)结果在state_dict()里看到一个奇怪的key还出现在buffers中。这是nn.Module的另一个默认行为普通tensor赋值给self.xxx会被自动当作buffer注册。如果你确实想保存这个状态这没问题但如果你只是想在forward里用一下临时tensor就不该挂到self上面应该先创建后传入或者用torch.zeros这种普通局部变量否则state_dict会越来越臃肿。5.2 用hook定位哪一层输出变成了NaN模型训练到一半loss突然变成NaN是每个人都遇到过的经典问题。我有一套固定打法先用hook在所有层上注册检查函数一旦发现输出中有NaN立刻打印层的名字、输入输出统计量并终止。def nan_check_fn(module, input, output): if torch.isnan(output).any(): print(fNaN detected in {module}) raise SystemExit handles [] for name, submodule in model.named_modules(): if isinstance(submodule, nn.Module) and hasattr(submodule, forward): handles.append(submodule.register_forward_hook(nan_check_fn))跑一次后第一个触发NaN的就是问题层。接下来再检查它的输入看是上一层的输出问题还是这一层自身的权重问题。这种排查方式不需要改forward里的任何代码通常几分钟就能定位到具体位置。需要注意的是named_modules()会递归遍历整棵模型树对每个子模块注册hook可能让前向变慢。但作为定位手段这点开销完全值得。定位完成后记得把handles全部remove()否则后续训练一直开着检查。5.3 让模型代码持续可维护的几个习惯最后分享几个我长期写PyTorch模型的习惯这些习惯都是从踩坑里提炼出来的第一个习惯所有层都在__init__里装配完成forward里只做逻辑编排。这样model.parameters()、state_dict()、to(device)、train()/eval()这些机制都不需要额外照顾所有状态自动一致。第二个习惯统一参数初始化。在__init__最后写一个self.apply(self._init_weights)把初始化逻辑集中在专门的函数里利用apply递归调用到所有子模块上的特性来一遍初始化def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.kaiming_uniform_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias)第三个习惯写完一个新模型先打印三样东西模型结构、state_dict()的key列表、以及喂一个随机tensor会不会正常forward和backward。前两步是验证装好了第三步是验证能跑通。我靠这个简单流程在把模型交付给训练脚本之前几乎可以过滤掉90%的常见问题。第四个习惯谨慎使用Python原生的list和dict存模块。在模型定义这个场景里原生容器只有在不涉及参数管理时才安全只要容器里的东西是nn.Module或nn.Parameter一律用ModuleList、ModuleDict、nn.Sequential。说回nn.Module官方文档其实很短但它设计的精妙程度远超过表面上那几百行代码。把赋值、收集、递归这些基础机制弄清楚之后你会发现模型创建这件事没有太多魔法可言——它就是一套严格的登记和递归规则。自己动手写一个自定义层再试着用hook查一次NaN出现的位置这套实战走下来比背十遍文档印象都深。