有一次我在排查一个模型导出问题发现在model.state_dict()里多出了几个既不是卷积核也不是偏置的张量加载权重的时候还因此报出了 missing keys。查来查去才意识到这段代码里的自定义模块特意用register_buffer()注册了一些均值向量。而要把这类“不参与梯度更新、却要跟着模型走”的张量按名字批量抓出来PyTorch 给的标准入口就是torch.nn.Module.named_buffers(prefix, recurseTrue)。这个 API 乍一看很简单但我在实际项目里用了不少次之后发现很多人对它的返回值、两个参数的边界行为以及它和state_dict、named_parameters的关联都理解得比较模糊。这篇就结合我自己的使用经验把 buffer 的本质、named_buffers的运行逻辑、源码里的隐藏机制以及几个真实踩坑场景一次讲清楚。1. 先搞清楚 buffer 是什么模型里不止有参数1.1 从 BatchNorm 的 running_mean 说起做深度学习的人几乎天天跟 BatchNorm 打交道但未必认真想过running_mean和running_var到底是怎么被管理起来的。训练时BatchNorm 根据当前 batch 统计出临时均值和方差做归一化同时还会用一个滑动平均更新一组长期统计量供推理时使用。这两组值有本质区别临时统计量每次前向都重新计算不需要保存长期统计量虽然参与前向计算却不该通过梯度下降去更新它只是被训练过程“顺带维护”的状态。问题来了如果不采取额外机制而只是把running_mean直接写成模块的一个 Python 属性那么模型调用.to(cuda)的时候这个张量不会自动跟着迁移model.state_dict()也不会包含它加载权重自然无从谈起。PyTorch 的解法就是 buffer 机制——通过register_buffer()注册的张量会被纳入模块的状态管理随着模块在不同设备之间迁移、随state_dict保存和恢复但不会出现在parameters()里也不会被优化器更新。理解这一点之后再看named_buffers的定位就很清晰了它是专门用来遍历模块树中所有这类“带名字的状态张量”的 API和named_parameters形成了两个互补的视角。1.2 buffer 与 parameter 到底差在哪为了后面讨论能建立在共识上我把两者的核心差异整理成一张表对比维度nn.Parameter参数buffer缓冲张量是否参与梯度计算作为可训练叶子节点不直接要求梯度requires_grad默认为 True默认为 False是否被优化器更新是优化器默认遍历它否靠自定义逻辑维护是否进入state_dict是默认是persistentTrue是否进入named_parameters是否是否进入named_buffers否是是否跟随.to(device)移动是是是否参与模型保存/加载是默认是需要特别提醒的是有些刚入门的朋友会把“requires_gradFalse 的张量”直接当成参数用比如在自定义模块里写self.fixed_weight torch.randn(...)然后参与前向计算。这个写法在 PyTorch 的Module.__setattr__规则下会被解释为注册了一个未命名的普通 buffer——因为你赋给属性的值是一个torch.Tensor且它既不是nn.Parameter也不是纯 Python 对象此时模块会把该 Tensor 自动注册到_buffers字典。所以它是“有名字的”但不会出现在named_parameters()里。这个小细节非常容易埋坑后面第 5 节我会专门展开。2. named_buffers 返回值与两个参数的细节行为2.1 返回的是什么一个可迭代对象以及它的命名规则named_buffers()的签名是named_buffers(prefix, recurseTrue)返回值是一个迭代器每次迭代产出一个(name, buffer)二元组。name 是字符串buffer 是torch.Tensor也允许是None延迟初始化场景。在最简单的用法里你可以直接把它转换成字典方便按名字访问import torch import torch.nn as nn class DemoModule(nn.Module): def __init__(self): super().__init__() self.register_buffer(mean_state, torch.zeros(4)) self.register_buffer(var_state, torch.ones(4)) self.conv nn.Conv2d(3, 8, 3) self.bn nn.BatchNorm2d(8) model DemoModule() buffers_dict dict(model.named_buffers()) for name, buf in buffers_dict.items(): print(name, buf.shape)运行这段代码输出大概是mean_state torch.Size([4]) var_state torch.Size([4]) conv.weight # 注意这是参数不是 buffer不会出现在上面 conv.bias bn.weight bn.bias bn.running_mean torch.Size([8]) bn.running_var torch.Size([8]) bn.num_batches_tracked torch.Size([])你会发现返回的名称使用了点分路径风格比如bn.running_mean这和named_modules()、named_parameters()的命名策略完全一致。这种命名方式在实际代码里特别方便因为你可以直接通过字符串后缀判断这个 buffer 属于哪种子模块然后做批量处理。一个容易被忽略的点是遍历顺序named_buffers返回的顺序基本遵循注册顺序也就是_buffers字典的插入顺序。如果模块内部注册顺序是mean_state先、bn相关的 buffer 后那么输出顺序也会是这个顺序。这意味着你不能假设遍历结果按字母序排列依赖顺序的代码最好先显式排序。2.2 prefix 参数的实际用途从使用角度讲prefix就是简单地给所有返回的名字前面拼一段字符串。我见过很多人在需要给不同模型加命名空间时会直接用这个参数比如同时处理 teacher 和 student 两组状态时teacher_buffers dict(teacher_model.named_buffers(prefixteacher)) student_buffers dict(student_model.named_buffers(prefixstudent))这样得到的 key 分别变成teacher.bn.running_mean和student.bn.running_mean在写保存逻辑或做蒸馏算法里的统计匹配时不会发生名字碰撞。另一个不那么直观但很实用的场景是在自定义容器模块里“模拟”父级路径。PyTorch 中子模块的 buffer 命名会自动带上前缀但如果你不想真正挂到某个nn.Sequential或者nn.ModuleList里只希望导出和查询时名字风格保持一致prefix就帮了大忙。此外在按前缀做过滤逻辑时统一加一段前缀再配合字符串匹配会比逐层判断模块类型更省心。2.3 recurseFalse 会让遍历边界发生什么变化recurseTrue是默认行为表示递归遍历所有子模块中注册的 bufferrecurseFalse则只返回当前模块自身直接注册的 buffer不去管任何子模块。这个参数在两类场景下非常有用。第一类是当你想严格区分“这个模块骨架自身维护的状态”和“内部所有子模块状态”时。比如一个封装好的残差块自身注册了一个shortcut_scale的缩放向量但你不想在更新这个缩放向量时把内部 BatchNorm 的running_var一起捞出来。此时用block.named_buffers(recurseFalse)就能精确拿到这一层自己拥有的状态。第二类是配合递归逻辑做定制遍历。默认的递归在有些复杂网络里不够用你需要自己控制子模块访问顺序或者只想处理特定类型的子模块。这时可以把recurseFalse跟手写递归结合先拿到当前模块的 buffer再按自己的策略遍历named_modules()并逐个调named_buffers(recurseFalse)。这种方式虽然麻烦一点但可控性最强尤其是在处理动态构建的子模块树时。3. 源码视角named_buffers 背靠的 _named_members 机制3.1 层层调用的调用链要真正吃透一个工具 API光看文档是不够的我一般会直接翻源码。named_buffers()的实现逻辑并不复杂本质上它是_named_members的一个特化版本def named_buffers(self, prefix, recurseTrue): return self._named_members( lambda module: ((name, module._buffers[name]) for name in module._buffers), prefixprefix, recurserecurse, )也就是说真正控制遍历的核心在_named_members。它的思路是如果recurseFalse当前待遍历的模块列表就是[(prefix, self)]。如果recurseTrue则先拿到self.named_modules(prefixprefix)得到所有后代模块及其名称前缀。对每个模块调用传入的get_members_fn去读取这个模块的_buffers字典。过滤掉None和已经出现过的 buffer 对象再拼出完整名字后yield。这里有个关键点get_members_fn遍历的是_buffers字典的键而不是所有属性。所以只要你用register_buffer注册过或者在某些赋值行为下被模块自动登记进_buffers它一定会出现在遍历结果里。反过来说如果某个张量只是普通属性比如self._tmp_tensor torch.zeros(2)它既不会进state_dict也不会被named_buffers返回因为它压根没进入_buffers字典。理解这个调用链的好处是当你需要扩展或复刻类似行为时可以直接复用_named_members的思路而不是重新迭代模块树。3.2 去重背后共享 buffer 的陷阱如果你仔细看_named_members的实现会发现它内部维护了一个memo集合用来记录已经返回过的 buffer 对象。一旦某个 buffer 对象已经出现过后续再遇到同名或同对象都会跳过。这个设计对应的是模型权重共享场景。举个例子我当年写一个孪生网络时两个上下游分支共享了同一个归一化统计量。常规情况下named_buffers会把同一个张量对象在两个分支下各返回一次但如果这个共享是通过“赋值同一个 Tensor 给两个模块属性”实现的PyTorch 的去重机制会导致只返回第一个分支下的名字。这个行为有时候是好事因为你确实不希望同一个对象被重复处理两次但有时候也会让人困惑——你想统计所有子模块的 buffer 数量结果数字比预期少。更麻烦的是named_buffers()的签名里根本没有remove_duplicate参数你没法直接关掉去重。如果你真的需要知道每个子模块各自拥有哪些 buffer就得自己手动循环for module_name, child in model.named_modules(): for buf_name, buf in child.named_buffers(recurseFalse): full_name f{module_name}.{buf_name} if module_name else buf_name print(full_name, buf.shape)这个写法不受共享去重影响因为它是以模块为维度逐层扫描的。代价是可能出现同一个底层张量被打印多次但如果你要的正是这种“按物理位置统计”的视角它反而更合适。3.3 persistentFalse 的 buffer 会不会出现在结果里PyTorch 从 1.6 开始支持register_buffer(..., persistentFalse)。很多人误以为persistentFalse表示“不持久因此 named_buffers 里也看不到”其实不是这样。我查过源码_buffers字典并不区分 persistent 与否persistent参数只影响state_dict的返回值。_named_members在遍历时只检查 buffer 是否为None以及是否重复不检查 persistent 标记。所以结果是非持久化的 buffer会出现在named_buffers()中但不会出现在state_dict()里。我在一个在线训练服务里就是这样用的把一些用于统计诊断的运行均值注册成persistentFalse的 buffer这样既可以在代码里按名字访问model.state_dict()里又干净整洁。如果你依赖named_buffers做全量状态扫描必须记住这一点否则会漏算非持久化 buffer反过来如果你在加载外部模型权重时用named_buffers做 key 匹配也会发现有些 buffer 是“只存在于内存、不存在于权重文件”的并不是模型文件出错了。4. 实战在项目中用 named_buffers 快速批量处理模型状态4.1 批量重置 BatchNorm 统计量训练过程中我经常遇到一个需求在做模型 warmup 或者 domain adaptation 实验时需要把模型里所有 BatchNorm 的running_mean和running_var重置成初始状态。逐个模块去找很蠢用named_buffers就可以几行代码搞定def reset_bn_stats(model): for name, buf in model.named_buffers(): if name.endswith(running_mean): buf.zero_() elif name.endswith(running_var): buf.fill_(1.0)这段代码的关键在于我没有依赖模块类型判断而是直接用名字后缀过滤。好处是即使未来换成别的归一化层只要它还维护running_mean之类的 buffer这段逻辑依然有效。缺点是如果某个自研模块恰好也叫这个名字会被误伤所以严谨一点的做法是结合模块名路径里的特征段来过滤。有人可能会问为什么不用model.apply(...)apply是按模块递归执行的虽然也能做但需要你写一个闭包而且到了模块内部还是要找 buffer。直接用named_buffers更扁平也更贴近“状态集中管理”的思路。4.2 按名称前缀筛选自研模块的 buffer我维护过一个自研的统计记忆模块它在每个网络层级都注册了几个原始统计量名字都有统一前缀比如mem_mean、mem_std。在写蒸馏脚本时我需要把所有这类状态一次性加载到另一组模块里。当时同事的写法是维护一个巨大的手工映射表每个层级写一行赋值代码。我接手后用named_buffers把这件事压缩成了一个基于前缀的筛选逻辑old_buffers dict(model.named_buffers()) new_buffers dict(new_model.named_buffers()) for key, value in old_buffers.items(): if mem_ in key: if key in new_buffers and new_buffers[key].shape value.shape: new_buffers[key].copy_(value.detach())这套方法同样适用于“只更新某类 buffer”的场景。比如在一个基于脉冲信号的工作里只需要将某一层之后的统计量重设我可以先用字符串解析拿到模块深度再做过滤。相比逐个模块手工处理这种做法更不容易漏。4.3 与 state_dict 配合实现 buffer 的校验与导出named_buffers还能用来做加载权重时的“buffer 差异体检”。我经常遇到的情况是换了新版本模型后旧的 checkpoint 里 buffer 数量对不上直接load_state_dict(strictTrue)会报错。这时与其看一屏错误信息不如自己先 diff 一遍def diff_buffer_keys(model, checkpoint): model_buffers set(dict(model.named_buffers()).keys()) checkpoint_buffers set(checkpoint.keys()) missing model_buffers - checkpoint_buffers unexpected checkpoint_buffers - model_buffers return missing, unexpected画重点这里checkpoint.keys()是包含参数和 buffer 的完整 key 集合所以直接和named_buffers()取并集比对不一定完整你需要先过滤出带running或自定义 buffer 后缀的 key。不过如果只是想快速知道差异规模这个简化版本的 diff 已经足够粗筛。我在做 ONNX 导出前也会跑一遍dict(model.named_buffers())确认所有 buffer 的形状和数据类型都是预期值。因为在导出过程中这些 buffer 会被固化为一组常量节点如果某个 buffer 是 fp64 而模型输入是 fp32导出的图里就可能出现隐式类型转换性能会莫名其妙地下降不少。5. 踩过的几个坑注册了却找不到找到了却不该改5.1 赋值 Tensor 却意外注册成 buffer在模块的__init__里直接写self.some_tensor torch.zeros(3)PyTorch 会把这个 Tensor 放到_buffers里而不会抛出任何提示。这看起来很方便但也是个雷。因为如果你本意是放一个临时变量比如用来记录一次前向里的中间结果那么这个“临时变量”就会跟着state_dict保存、跟着.to()迁移也会出现在named_buffers()里。我遇到过最讽刺的 bug 是训练脚本里有个模块错误地把一个庞大的临时解码矩阵值赋给了普通属性结果每个 checkpoint 竟然膨胀了几百 MB而排查了很久才发现它被静默注册成了 buffer。从那以后我在自定义模块里凡是写属性赋值都会先分辨这个属性是 Tensor 吗是的话它到底应不应该作为 buffer 存在如果是就用register_buffer显式注册并想清楚是否要设为persistentFalse如果只是临时量就把属性名加下划线前缀并且在赋值前用类似self._tmp None的方式规避模块的__setattr__路径。这个习惯可以帮你省下大量排查时间。5.2 None 占位 buffer 与延迟初始化register_buffer允许注册一个值为None的占位 buffer。这是为了满足延迟初始化需求网络结构先确定好名字具体形状等第一次 forward 看到输入数据后再决定。但_named_members在遍历时会跳过值为None的条目——这本身是正确行为因为一个没有实际张量的 buffer 不该被当成有效状态处理。可这里有个很隐蔽的坑如果你的代码依赖named_buffers返回结果去初始化某些运行状态而这些 buffer 恰好还没有被填充得到的是一个空集合。反之如果在遍历之后补上了None的值named_buffers的结果又变了。这就导致“先判断模型是否有某些 buffer再决定走哪条初始化逻辑”的代码可能在训练的不同阶段给出不一样的结论。我的做法是在模型构建完成后主动触发一次 dummy forward把所有延迟 buffer 填充好再进入后续流程。这样初期看着有点浪费但能保证任何依赖named_buffers的代码看到的都是一致、完整的视图。5.3 多个父模块共享同一个 buffer 时被去重“吞掉”前面讲源码时已经说过去重机制。实际业务中最典型的情况是老师和学生模型共享同一个统计模块或者多个子模块被同一个nn.Module实例反复引用。这时named_buffers()的输出里只会出现第一次遇到的位置。如果你只是做展示或统计这没问题但如果你写了一个自动匹配脚本希望把所有出现该 buffer 的模块位置都更新一遍就会漏掉一部分。我踩过一次后总结了检查方法先用named_buffers()全量扫一遍再用 5.2 节的手动逐模块扫描方式扫一遍对比两次数量差异。如果发现差异就去检查模块树里是不是有共享引用。遇到共享引用时要么接受“同一状态只维护一份”的语义要么在代码层面显式保证所有使用位置都指向同一个更新函数。6. 扩展想法从 named_buffers 延伸到自定义状态管理6.1 自定义模块中如何规划 buffer 命名使用named_buffers越频繁越会发现命名规范的重要性。它不只是为了调试好看更是为了后续批量操作的可行性。我个人会遵循几条原则所有自定义 buffer 统一放在一个子命名空间里比如用aux_前缀避免和内置 BatchNorm 或 LayerNorm 的running_*混淆。一个模块内 buffer 尽量用动词性含义开头比如accum_buffer、counter_tensor方便一眼看出它在做什么。如果存在明显的一对状态比如均值和方差最好统一后缀比如_mean、_var这样遍历后可以用字符串 replace 快速成对处理。这些原则看起来像审美问题但在写大规模训练框架时它直接决定了你能不能用“前缀匹配 named_buffers”这套组合拳快速实现通用逻辑。6.2 用 named_buffers 实现 EMA 均值池化有一类非常常见的需求是维护指数滑动平均比如模型参数的 EMA或者某些 intermediate feature 的均值。把 EMA 状态注册为 buffer 是最稳妥的做法因为它天然会随模块移动、随 checkpoint 保存。下面是一个简化但可运行的示例class EMAState(nn.Module): def __init__(self, dim, decay0.99): super().__init__() self.decay decay self.register_buffer(running_ema, torch.zeros(dim)) torch.no_grad() def update(self, new_value): self.running_ema.mul_(self.decay).add_(new_value, alpha1 - self.decay) def get_ema(self): return self.running_ema有了这个模块后你要想在整个网络里统一把这一类 EMA 状态都禁用或清零就不需要知道它具体挂在哪一层直接遍历named_buffers()for name, buf in model.named_buffers(): if name.endswith(running_ema): buf.zero_()这里隐含着一个容易犯的错误更新 buffer 时一定要包在torch.no_grad()环境里或者使用原地方法。因为一旦你在有梯度的上下文里对 buffer 做非原地操作计算图就会把 buffer 当作叶子节点追踪后续反向传播可能得到错误梯度。虽然 buffer 本身requires_gradFalse但它如果作为某个运算的输入梯度照样会流向它参与的那条路径这种“意外参与反向传播”的 buffer 是训练异常的常见来源之一。6.3 检查点迁移时常见的 buffer 不匹配问题最后聊聊跨版本加载模型的场景。模型的权重文件里既有参数也有 buffer而 buffer 的数量和形状在不同版本之间经常发生变化。比如旧版 BatchNorm 里num_batches_tracked还没有引入或者新版自定义模块新增了一个running_emabuffer。直接load_state_dict会给你一堆 unexpected keys 和 missing keys非常劝退。我当时写了一个自动对齐函数原理就是利用named_buffers做 key 集合求差def load_with_buffer_compat(model, checkpoint, default_fill0.0): state model.state_dict() for key, buffer_value in model.named_buffers(): if key not in checkpoint: state[key] torch.full_like(buffer_value, default_fill) else: ckpt_value checkpoint[key] if ckpt_value.shape buffer_value.shape: state[key] ckpt_value else: state[key] torch.full_like(buffer_value, default_fill) model.load_state_dict(state, strictFalse)这套兼容逻辑解决了一个非常现实的问题新加一个统计量 buffer 之后旧 checkpoint 没有对应值如果硬加载会导致新模型从零开始积累这个统计量效果在初始阶段会有点偏差。用默认值填充至少能保证前向过程不崩溃至于什么时候该清零、什么时候该填充默认统计值就需要结合业务来判断了。我试过一个小技巧把缺失的 buffer 默认值打印出来在前几次迭代里监控它的变化曲线基本能判断它是否需要在一开始就重新统计。比如running_mean如果默认是 0训练后慢慢贴近真实分布是正常的但如果是某些尺度敏感的自定义状态可能就需要在加载时手工复制某个相近模块的值而不是盲目填 0。如果你也在维护一个长期迭代的模型项目我的建议是把“用named_buffers做状态检查”写成一个标准的调试工具函数挂在训练脚本的启动阶段每次加载权重后自动打印一遍所有 buffer 的(name, shape, dtype, device)信息。这个习惯在排查模型导出、蒸馏、剪枝、量化、多卡加载等问题时会帮你省下大量时间。
