如果只用PyTorch训练过几个现成网络很少会认真琢磨nn.Module到底是怎么把参数管起来的。但一旦开始自己写模型、改结构、做迁移学习就会发现所有莫名其妙的问题——参数没更新、设备不对、保存后加载失败——十有八九都出在对模型构造的理解上。这篇就围绕torch.nn.Module把模型构造这件事讲透从它解决的问题、底层的注册机制到自定义模块的完整套路再到踩坑实录看完基本能摸清PyTorch模型的骨架结构。这篇适合刚学完PyTorch基础、想自己动手写模型的同学也适合写了不少网络但遇到疑难杂症的老手。理解Module的核心不光是会用class Net(nn.Module)这种模板而是真正搞懂Python类和PyTorch框架之间到底在哪一层发生了什么事。1. 为什么要有Module先看没有它时有多痛苦1.1 手动管理参数的原始做法回到最基础的场景实现一个简单的线性层y Wx b。如果完全不用nn.Module代码大概是这样的import torch # 手动创建参数 W torch.randn(4, 3, requires_gradTrue) b torch.zeros(4, requires_gradTrue) def linear(x): return torch.matmul(x, W.t()) b def update_params(lr): with torch.no_grad(): W - lr * W.grad b - lr * b.grad W.grad.zero_() b.grad.zero_()单看这个例子似乎也不算麻烦。但把场景放大一个10层的网络每层有自己的权重和偏置还要考虑每层用不用dropout、要不要BatchNorm、不同层可能需要不同的初始化方式这时候手动管理参数就完全失控了。你要自己维护一个巨大的列表记录每组的顺序、shape、设备位置还要小心model.cuda()时有没有漏掉哪一组。更关键的是这种写法里的W和b只是恰好是一对的独立张量。框架没办法自动推导它们的关系也没法自动做梯度回传、参数保存、设备迁移。整个训练流程里你都得手动兜底。1.2 Module到底做了什么抽象nn.Module做的核心事情可以概括为一句话把参数、子模块、前向传播逻辑绑定成一个整体并提供一套统一管理机制。它给我的感觉很像一个项目的负责人——负责登记所有成员参数、管辖所有小组子模块、规定干活流程forward还要统一处理考勤汇报state_dict 和设备迁移。具体来说它提供了几个很关键的能力自动登记参数所有被赋值为nn.Parameter的属性都会自动进入parameters()迭代器自动登记子模块所有被赋值为nn.Module的属性都会自动成为子模块递归展开统一设备迁移调用model.cuda()时所有参数和缓冲区一并迁移统一的序列化state_dict()把全模型的关键数据打包成有序字典统一的训练/评估切换model.train()和model.eval()控制全模型状态这些能力不是魔法而是通过Python的魔术方法实现的。下一节讲这里面的门道。2. Module的底层机制赋值时的拦截逻辑2.1__setattr__如何决定一个属性的命运Python类里给实例赋值属性本质上是调用__setattr__方法。Module恰恰重写了这个方法所以它对你往里面塞了什么极其敏感。这也是理解Module最核心的一层你在__init__里写的每一行赋值语句都会经过这个拦截逻辑分类处理。大致规则是这样的被赋值的对象类型去向举例torch.nn.Parameter_parameters有序字典self.weight nn.Parameter(torch.randn(...))nn.Module的实例_modules有序字典self.fc nn.Linear(10, 5)普通张量Tensor普通属性字典不参与统一管理self.some_tensor torch.zeros(...)其他任意对象普通属性self.dropout_rate 0.5其中前两类是Module最关注的对象因为PyTorch要求参数和子模块在注册后的生命周期内能被框架统一发现、统一操作。可以自己验证一下import torch import torch.nn as nn class MyNet(nn.Module): def __init__(self): super().__init__() self.w nn.Parameter(torch.randn(3, 4)) self.layer nn.Linear(4, 2) self.not_managed torch.randn(2, 2) # 这个不会出现在 parameters() 里 net MyNet() print(net._parameters) # OrderedDict([(w, Parameter(...))]) print(net._modules) # OrderedDict([(layer, Linear(...))])这段代码里有个很关键的细节没有体现——self.w被赋值时内部还有一个显式调用register_parameter(w, ...)的过程。你可以手动注册参数但日常用法里nn.Parameter赋值就足够了。2.2 为什么要用OrderedDict来存OrderedDict是Python中能保持插入顺序的字典Module用OrderedDict而不是普通字典来存储_parameters和_modules原因在于顺序就是网络结构。比如网络先是一个卷积层再是一个线性层。当调用model.state_dict()时键名会按照layer1.weight、layer1.bias、layer2.weight、layer2.bias这样的顺序排列。保存和加载时严格对应。如果换成普通字典Python 3.7虽然也保留了插入顺序但OrderedDict在这里更明确地传递了顺序重要的语义。更深一层Module在递归展开子模块时必须保持从外层到内层、从前到后的确定性顺序才能和state_dict的键一一对应。这点在断点续训、跨设备部署时尤其重要——加载参数时所有张量必须按名字逐一对应上顺序错乱会导致严重问题。2.3 参数注册与普通Tensor的区别初学者最容易疑惑的问题self.weight nn.Parameter(torch.randn(...))和self.weight torch.randn(...)有什么区别两者都是给实例赋值但命运完全不同。nn.Parameter是Tensor的子类。Module看到Parameter就会自动注册到_parameters。而普通Tensor被赋值后只是普通属性不会出现在parameters()里也不会出现在state_dict()中更不会跟着model.cuda()迁移。换句话说如果你把需要训练的张量用普通Tensor赋值模型可以正常跑因为Python层面self.weight还是同一个张量前向传播也能用。但训练时优化器只会拿model.parameters()去更新普通张量根本不在列表里于是你会看到loss一直不下降。这种问题特别隐蔽。另外如果某个参数不需要参与训练但需要随模型保存比如BatchNorm的running_mean和running_var应该用register_buffer()注册为缓冲区。缓冲区不进优化器但会进state_dict、会随设备迁移。这是自写归一化层时的必修课。3. 自定义Module的实操套路从零写一个可训练层3.1 初始化参数创建与绑定自己写一个简单的加权求和层比只调现成nn.Linear更能理解参数管理的流程。下面这个例子实现了一个可训练的对输入每个特征做加权缩放的自定义层import torch import torch.nn as nn class FeatureScale(nn.Module): def __init__(self, in_features, init_scale1.0): super().__init__() self.scale nn.Parameter(torch.full((in_features,), init_scale)) self.bias nn.Parameter(torch.zeros(in_features)) # 用 init 模块做更精细的初始化 nn.init.uniform_(self.scale, 0.5, 1.5) def forward(self, x): return x * self.scale self.bias注意__init__里的顺序先调用super().__init__()再创建参数。原因在于父类的Module.__init__会初始化那两个有序字典和各类状态标志。如果先赋值参数再调用super().__init__()赋值时_parameters还不存在就会抛异常。这一步在网上很多教程里被一笔带过但实际开发中一旦漏掉就是程序直接崩溃的下场。参数创建后建议用torch.nn.init里的函数做初始化。为什么不能直接torch.randn赋值倒不是说不行而是nn.init统一封装了常见的初始化策略语义更清晰而且像xavier_uniform_、kaiming_normal_这类初始化方式写起来更方便、更容易统一管理。3.2 前向传播定义的关键习惯forward是Module的子类必须实现的方法。但需要特别注意的是forward不是直接被外部调用的常规调用方式是model(x)。Module重写了__call__方法__call__内部会执行一些钩子和前置逻辑这段可能是forward_pre_hooks和forward_hooks然后才进入forward本身。所以如果你在外部写model.forward(x)其实也可以跑但会绕过钩子机制。训练时nn.Module本身不强制检查这一点但在用nn.DataParallel或某些自定义 trainer 时钩子不执行会影响预期行为。标准做法就是永远通过model(x)做前向。再提一个实操细节forward里尽量用torch函数而不是Python原生逻辑。不是说不能用if、for而是所有重计算部分得是可微张量运算。比如特征筛选、条件分支、循环这些都没问题但如果你把张量转成numpy再转回来梯度就断了模型完全无法训练。3.3 缓冲区那些需要保存但不需要梯度的量很多自写层不只是有可训练参数还要维护一些运行时统计量。最典型的就是BatchNorm里的running_mean和running_var它们会随着前向传播不断更新但更新方式不是通过优化器梯度下降而是通过momentum做的滚动平均。这时候就该用register_bufferclass MyNorm(nn.Module): def __init__(self, num_features, momentum0.9): super().__init__() self.momentum momentum self.register_buffer(running_mean, torch.zeros(num_features)) self.register_buffer(running_var, torch.ones(num_features))注册为buffer之后model.cuda()会一并迁移、model.state_dict()会保存这两个量加载模型时也不会丢失。而且它们不会出现在优化器的参数列表里也就不会产生梯度更新相关的误操作。buffer还有一个细节model.to(device)时会自动迁移但如果你直接用self.running_mean self.running_mean.cuda()这种方式重新赋值由于它是普通属性赋值覆盖了原来注册的位置反而可能破坏注册关系。正确做法是始终调用model.to(device)或model.cuda()统一处理。4. 子模块装配与组合Sequential、ModuleList、ModuleDict怎么选4.1 三种容器背后的设计考量PyTorch提供了几种容器类来组织子模块很多人只是随手用但对它们的适用场景并不清楚。这里我把它们放在一起对比容器底层结构前向方式适用场景nn.Sequential_modules有序字典按顺序自动执行简单堆叠、无分支的网络nn.ModuleList类似列表的_modules存储需要手动遍历调用数量不定的同构子模块nn.ModuleDict类似字典的_modules存储需要按键取值调用按名称动态选择分支nn.Sequential最省事因为它自带一个简化版forward会把传入的所有层按顺序串联执行。适合那种输入进来一路往下走的网络比如经典的卷积堆叠后接全连接self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2) )但Sequential也有明显局限它没有分支能力。如果前向里需要把某一层的输出同时送到两条路径或者需要跨层相加Sequential就力不从心了。此时应该换用ModuleList或直接的子模块赋值方式。4.2 嵌套结构与模型自由度ModuleList和ModuleDict更适合结构灵活的场景。比如自己实现一个多分支网络分支数量在初始化时不确定或者需要在forward里根据输入动态选择某些层ModuleList是更好的选择class MultiHeadNet(nn.Module): def __init__(self, hidden_dim64, num_heads3): super().__init__() self.backbone nn.Sequential( nn.Linear(128, hidden_dim), nn.ReLU(inplaceTrue) ) self.heads nn.ModuleList([ nn.Linear(hidden_dim, 10) for _ in range(num_heads) ]) def forward(self, x, head_idxNone): feat self.backbone(x) if head_idx is not None: return self.heads[head_idx](feat) # 默认返回所有头的输出 return [head(feat) for head in self.heads]注意这里用的是nn.ModuleList而不是Python的list。如果用原生的list包住这些Linear层它们不会被_modules登记model.cuda()不会迁移、model.parameters()不会包含它们模型保存时也会丢失这些层的参数。这个坑在第6章还会详细讲。ModuleDict最常见的场景是按名称切换不同的激活函数、不同的注意力头、不同的分类器。比如代码里有一个字符串参数指定用哪种注意力机制直接从这个字典里取对应的Module来调用逻辑非常清晰。4.3 嵌套结构的参数名规则嵌套模块的state_dict键名是从根到叶的路径中间用点号连接。例如一个模型里self.features下有self.conv1那么键名就是features.conv1.weight。理解这套规则对加载预训练模型、做模型手术、调试参数都很有帮助。一个小技巧调用model.named_parameters()不仅能看到参数还能看到参数所属的完整路径。排查哪个参数没被更新时先打印一遍这个列表比什么都直观。5. 参数管理的核心操作parameters、named_parameters与state_dict5.1 遍历参数做初始化自定义模型时经常需要批量初始化。Module.parameters()方法递归遍历所有已注册的参数配合nn.init可以写统一的初始化逻辑def init_weights(model): for name, param in model.named_parameters(): if weight in name: nn.init.kaiming_normal_(param, modefan_out, nonlinearityrelu) elif bias in name: nn.init.constant_(param, 0)这里的name是形如features.conv1.weight的完整路径。如果你想只对某些层做特殊初始化用name做前缀匹配非常方便。这种按名字筛选的方式比遍历所有参数再猜类型要可靠得多。5.2 冻结参数实现迁移学习迁移学习里常见操作是冻结骨干网络的参数只训练分类头。实现方式很简单把需要冻结的参数的requires_grad设为Falsefor name, param in model.named_parameters(): if name.startswith(backbone): param.requires_grad False optimizer torch.optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr1e-3 )这里filter是关键即使你把参数的requires_grad设成Falsemodel.parameters()仍会返回它。如果不做过滤优化器还是会拿它去尝试更新虽然梯度在反向传播时不会为它计算但也会浪费显存、产生奇怪警告。稳妥做法就是在构建优化器时过滤掉不训练的参数。还有一个隐含问题冻结某层参数后它的子模块依然处于model.train()的训练模式如果冻结的是带BatchNorm的层running_mean仍会被更新。要彻底冻结还需要把对应子模块设为eval()模式。5.3 保存与加载模型的状态字典保存和加载模型是所有构造问题中操作最频繁的。正确姿势是只保存state_dict而不是整个模型对象# 保存 torch.save(model.state_dict(), model_weights.pth) # 加载 model MyModel() model.load_state_dict(torch.load(model_weights.pth, map_locationcpu))为什么不推荐保存整个模型对象因为torch.save(model)会把模型的类定义路径也打包进去加载时要求环境里必须有同名同路径的类。一旦项目目录结构变了或者类名改了就会加载失败。只保存state_dict则灵活得多。load_state_dict默认strictTrue要求键名和当前模型的state_dict完全匹配。需要部分加载时可以先把预训练权重载入一个字典再手动过滤掉不匹配的键或者设置strictFalse。strictFalse会忽略缺失的键和多余的键适合做模型结构微调后的加载但一定要仔细看输出日志里报的缺失键否则容易静默丢掉某些层的参数。6. 踩坑实录与排查技巧6.1 Module里放Python list的坑这是网上出现频率最高的错误在__init__里写self.layers [nn.Linear(10, 10) for _ in range(3)]。肉眼看着没问题模型也能跑直到某天发现model.cuda()之后某些层还在CPU上或者parameters()里少了这些层的参数才意识到list里的Module根本没有被注册。修复方式很简单把原生list换成nn.ModuleList。如果你想在debug时确认结构可以打印model._modules看哪些子模块被登记了、哪些没有。这个检查在模型规模变大后尤其重要。6.2 把Tensor当成Parameter塞进Module另一种常见错误是把自定义层的权重写成了普通Tensor。比如self.weight torch.randn(in_features, out_features)由于普通Tensor不是ParameterModule不会把它登记到_parameters。它虽然作为一个属性可以正常访问但不会出现在model.parameters()里也就不会被优化器更新。训练时表现为loss在某个初始值附近震荡甚至完全不降很容易误判成学习率或初始化的问题。排查技巧打印model.parameters()的列表长度和预期对比如果少了参数基本就是这个原因。修复方式就是手动包一层nn.Parameter或者用register_parameter显式注册。6.3 设备不一致导致的运行时报错在多GPU或混合设备场景下模型的某一部分参数在CPU、另一部分在GPU前向传播到某个矩阵乘法时就会报Expected all tensors to be on the same device。排查思路是先统一调用model.cuda()或model.to(device)再遍历打印每个参数和设备for name, param in model.named_parameters(): if str(param.device) ! cuda:0: print(name, param.device)另外如果自定义forward里有临时创建的张量比如torch.zeros(...)、torch.arange(...)默认是在CPU上创建需要显式使用与输入相同的设备或者直接用torch.zeros_like(x)这类带like的函数来保持设备一致。这是很细节的部分但也是多设备训练中最常见的坑。6.4 eval模式与推理行为的混淆model.eval()和model.train()只影响BatchNorm、Dropout等有训练行为和推理行为差别的层它不会自动禁用梯度计算。推理时如果不主动包torch.no_grad()PyTorch仍会构建计算图、占用显存但结果在数值上不受影响。有些同学误以为调了model.eval()就不会再更新BatchNorm的统计量了。事实是eval()确实让BatchNorm使用缓存的running_mean/running_var所以前向不会改动缓冲区。但在train()模式下即使requires_gradFalseBatchNorm仍会根据当前batch更新统计量和参数是否冻结无关。实际操作中的一点体会把nn.Module的结构、注册机制、参数管理弄明白之后写模型就是按规范组织参数和子模块这么简单。遇到奇怪的训练问题先检查_parameters、_modules、buffers这三个内部字典再检查优化器拿到的参数列表多半能定位问题。我个人写自定义模型时习惯在__init__末尾加一句print(self)或者打印state_dict()的keys确认结构符合预期再开始训练。这个小动作帮我避免了很多因为参数没注册、层放错位置导致的隐性bug。模型构造是地基地基稳了后面所有训练技巧才能真正派上用场。