企业官网建设流程全解析

开头正题。前几天我在排查一个推理性能异常的问题，模型加载后输出结果错得离谱，逐层打印特征图才发现，BatchNorm 层的 running_mean 和 running_var 全是初值 0 和 1。问题不在模型的 forward 逻辑，而是某些张量压根没被当作模型的“状态”保存下来，训练完只存了参数，统计量全丢了。

这个坑背后牵出来的就是torch.nn.Module.buffers(recurse=True)这套机制。说实话，这个 API 看着简单，不就是把模块里的 buffer 拿遍历一遍吗？但很多人在state_dict里看到running_mean、running_var这类键时，并不完全清楚它们和weight、bias这些参数的区别，更不清楚buffers()在什么场景下能用出真正的价值。这篇文章就把这个概念彻底拆开，从“为什么需要 buffer”讲到“遍历、注册、保存、加载的完整实操”，最后再分享几个我实际踩过的坑。

1. 先搞清楚 buffers 到底是什么

1.1 模型里的张量，其实有三种身份

PyTorch 的nn.Module内部维护的不只是大家最熟悉的_parameters字典，还有一个_buffers字典。模块里的每个张量，按身份可以分成三类：

• Parameter：注册在_parameters里，requires_grad=True，会被优化器更新，也会进入state_dict。
• Buffer：注册在_buffers里，requires_grad=False，不参与梯度更新，但会进入state_dict。
• 普通 Tensor 属性：就是直接self.xxx = torch.tensor(...)贴在模块上的张量，三者都不属于，既不进state_dict，也不会跟着model.to(device)移动。

很多人对“模型参数”有个直觉误区，认为所有需要保存的张量都是 Parameter。但实际上像 BatchNorm 的running_mean和running_var、EMA（指数移动平均）里的影子权重、自回归模型里的状态缓存，这些张量需要随着模型一起保存和移动，却绝不能被梯度更新。如果把它们注册成 Parameter，优化器就可能会更新它们，还会让它们混入反向传播过程，造成不想要的行为。

Buffer 本质上是“模型需要随身携带的、但不参与梯度计算的状态”。可以把它类比成行李箱上的标签：它不是行李箱里装的内容本身，但它跟着行李箱走，丢了它，箱子就找不到归属。

所以，判断一个张量该不该用register_buffer注册，关键就看两个条件：

1. 它是否需要进入state_dict，实现保存和加载？
2. 它是否需要跟着模型在 CPU、GPU 之间迁移？
3. 它是否完全不需要梯度？

只要前两个答案是“是”，第三个答案也是“是”，那就该用 buffer。

1.2 一张表看清 Parameter 和 Buffer 的差异

这个表格里最容易忽略的是最后两行。model.parameters()遍历不到 buffer，这是刻意设计的，避免优化器把 buffer 当作可学习参数收集进去。而model.buffers()遍历不到普通 Tensor 属性，这一点会让很多人栽跟头，我后面在常见问题里会详细讲。

1.3 持久化 buffer 和非持久化 buffer

register_buffer()还有一个很容易被忽略的参数：persistent。这个参数在比较新的 PyTorch 版本里才暴露出来，语义非常清楚：

• persistent=True（默认）：buffer 会进入state_dict()，保存权重时会一并保存。
• persistent=False：buffer 不会进入state_dict()，但依然存在于_buffers字典中，依然会被model.buffers()遍历到，也依然会跟随设备移动。

非持久化 buffer 很适合保存“只需本次运行期间存在的临时状态”。比如某些层内部的计数、缓存、掩码，这些状态当前进程用得上，但不应写入 checkpoint，否则不同 device 的 checkpoint 结构会对不上。

我在实际项目里用过非持久化 buffer 保存一个“运行中模块实例的唯一 ID”，这个 ID 只用来调试时追踪模块实例，完全没有持久化价值。如果不用persistent=False，它会被写入 checkpoint，导致不同卡上保存的结构不一致，非常麻烦。

2. 把buffers()方法彻底拆解

2.1 最基本的用法和返回值

model.buffers(recurse=True)的返回值是一个生成器（Generator），迭代时逐个产出模块直接或间接持有的 buffer 张量。注意，它产出的是张量本身，而不是键值对。如果你想同时拿到名字和值，应该用named_buffers()。

import torch import torch.nn as nn class DemoModule(nn.Module): def __init__(self): super().__init__() self.bn = nn.BatchNorm2d(16) self.register_buffer("my_state", torch.zeros(4)) model = DemoModule() # buffers() 返回生成器，直接打印看不出东西 print(model.buffers()) # <generator object Module.buffers at 0x...> # 必须 list 或循环 for buf in model.buffers(): print(buf.shape)

因为返回的是生成器，每次调用都会重新遍历一次。别想着复用一次迭代结果，比如：

bufs = model.buffers() # 第一次遍历 for b in bufs: pass # 第二次遍历，什么都没了 for b in bufs: pass

这时就应该用list(model.buffers())把结果固化下来，或者干脆用named_buffers()。

2.2 recurse=True 到底递归了什么

recurse=True的意思是要不要递归遍历子模块的 buffer。默认就是True。遍历顺序遵循模块注册树的深度优先顺序，先遍历模块本身直接注册的 buffer，然后按子模块注册顺序，逐个递归遍历子模块的 buffer。

model = nn.Sequential( nn.BatchNorm1d(8), nn.Linear(8, 4), ) model.to("cuda" if torch.cuda.is_available() else "cpu") for name, buf in model.named_buffers(): print(name, buf.shape)

输出大致是：

0.running_mean torch.Size([8]) 0.running_var torch.Size([8]) 0.num_batches_tracked torch.Size([])

这里的0.running_mean，键就是由“子模块路径 + buffer 名”拼接而成。路径和名称之间用点号连接。

如果传recurse=False，只会返回当前模块_buffers字典里直接存放的 buffer，子模块里的一个都不返回。

还有一个细节，在较新的 PyTorch 版本中，named_buffers()可以额外设置remove_duplicate参数，默认是True。当多个父模块共享同一个子模块实例时，这个参数会决定 buffer 是否被重复返回。比如同一个bn = nn.BatchNorm2d(...)同时挂在两个 Sequential 里，设置remove_duplicate=False后就会看到同一个 buffer 出现两次。

2.3 与 state_dict、named_parameters 的相互补充

很多人搞不清model.named_buffers()和model.state_dict()的关系，其实它们是两个维度的东西：

state_dict()返回的是一个OrderedDict，其中参数的键和持久化 buffer 的键是混在一起的。比如：

state = model.state_dict() for k, v in state.items(): print(k, v.shape)

这个方法通常用作 checkpoint 的保存对象。而named_buffers()更适合在代码运行时对 buffer 做定向检查和修改，尤其适合遍历所有 BatchNorm 统一处理场景。

我个人的习惯是：需要序列化时用state_dict()，需要运行时操作张量时用named_buffers()（配合buffers()），需要喂给优化器时用parameters()。三者各管一段，不要混着用。

3. 三个必须用 buffers 的真实场景

3.1 BatchNorm 统计量的管理和同步

BatchNorm 是最常见的 buffer 载体。它的running_mean和running_var就是在训练过程中由 forward 函数内部更新的，更新方式是直接对张量做 inplace 操作：

running_mean = (1 - momentum) * running_mean + momentum * batch_mean

这类统计量完全满足“需要保存、需要跟随设备、不需要梯度”的三要素，所以它们天然就是 buffer。

实际工程里，我经常需要手动操作这些统计量，最典型的场景是模型剪枝或迁移学习时对 BN 层做统计量重校准：

def reset_bn_stats(model): for module in model.modules(): if isinstance(module, nn.BatchNorm2d): module.reset_running_stats()

这种全模型统一操作的场景，用named_buffers()会更灵活。比如想只看 running_mean 的值分布，可以：

means = [] for name, buf in model.named_buffers(): if name.endswith("running_mean"): means.append(buf.flatten())

在分布式训练里，buffer 还有一个容易被忽略的行为：很多框架（例如 DDP）默认会在 forward 前把 rank 0 的 buffer 广播到其他 rank，目的是确保 BatchNorm 统计量在所有进程间保持一致。正因为这些统计量是 buffer，框架才能自动识别并同步。如果当初被设计成普通 Tensor 属性，同步逻辑就不可能统一遍历到它们了。

3.2 滑动平均模型（EMA）的影子权重

EMA 是深度学习训练里非常常用的技巧，它会在训练过程中维护一份“参数的移动平均”，推理时用这份平滑后的权重，往往能得到更稳的结果。实现 EMA 时，最容易犯的错误就是把影子权重保存成普通 Tensor 或 Parameter。

正确做法是把影子状态注册为 buffer：

class ModelWithEMA(nn.Module): def __init__(self, backbone, decay=0.999): super().__init__() self.backbone = backbone self.decay = decay # 为每一个可训练参数创建影子参数 for name, param in backbone.named_parameters(): self.register_buffer( f"ema_{name}", param.detach().clone() )

这样影子权重会随着model.to(device)自动迁移，也会随着model.state_dict()保存下来。同时因为它们不参与反向传播，所以requires_grad保持为 False，天然不会被优化器跟踪。

训练循环里更新影子权重时，核心操作是：

@torch.no_grad() def update_ema(model): for name, param in model.backbone.named_parameters(): ema_buf = getattr(model, f"ema_{name}") ema_buf.mul_(model.decay).add_(param.detach(), alpha=1 - model.decay)

这里有个隐藏问题：影子权重的名字是通过f"ema_{name}"动态注册的，如果网络结构里有参数名包含点号点，那么注册时就会抛出错误。因为register_buffer不允许名字里包含点号。遇到这种情况，我会先把名字里的点号替换成下划线，再拼上前缀。这个坑我在复杂模型（比如某些检测头的嵌套模块）里真实遇到过。

3.3 自定义状态张量的标准姿势

除了 BN 和 EMA，还有一类场景也常用 buffer：模型内部需要维护统计量或定时器。

举个例子，我封装过一个 CNN 里的在线归一化层，它需要维护一个“当前 batch 的样本数”和一个“特征维度的累计和”。这两个状态都不需要梯度，但必须在训练过程中更新，并随 checkpoint 保存。于是：

class OnlineNormalization(nn.Module): def __init__(self, num_features): super().__init__() self.register_buffer("running_count", torch.zeros(())) self.register_buffer("running_sum", torch.zeros(num_features)) def forward(self, x): self.running_count += x.size(0) self.running_sum += x.sum(dim=(0, 2, 3)) if x.dim() == 4 else x.sum(0) return x

这里如果把running_count做成普通属性，模型保存时它不会被保存，而且把模型从 CPU 搬到 GPU 时它也不会跟着走，最后要么重新初始化，要么隐式出错。这几乎都是我在项目里亲眼见过的教训。

4. 实操：注册、遍历、修改与管理 buffer 的完整流程

4.1 注册一套自定义 buffer

register_buffer的完整签名大致是register_buffer(name, tensor, persistent=True)。从参数名能看出，它既有默认持久化的真脸，也有非持久化的一面。完整注册流程通常放在模块的__init__里，最晚也要在第一次使用前完成，因为state_dict的 key 结构是在模块构造和赋值过程中逐步确定下来的。

看看一个稍微完整一点的示例：

import torch import torch.nn as nn class CustomModel(nn.Module): def __init__(self): super().__init__() self.conv = nn.Conv2d(3, 16, 3, padding=1) # 持久化 buffer：会保存 self.register_buffer("num_steps", torch.tensor(0, dtype=torch.long)) # 非持久化 buffer：不保存 self.register_buffer("debug_mask", torch.ones(1), persistent=False) def forward(self, x): # 在 forward 中更新 buffer 需要在 no_grad 下 with torch.no_grad(): self.num_steps.add_(1) return self.conv(x) model = CustomModel() print(dict(model.named_buffers())) print(set(model.state_dict().keys())) # state_dict 里有 conv.weight、conv.bias、num_steps，没有 debug_mask

有几个关键细节：

第一，register_buffer不允许name里带点号，也不允许和已有属性重名。如果名字冲突，会抛出KeyError。

第二，如果传入的tensor是None，register_buffer会从_buffers里移除这个键。利用这个特性可以注销 buffer。

第三，buffers()遍历的是_buffers字典，非持久化 buffer 也会被遍历到。所以即使debug_mask不在state_dict里，它依然会被model.buffers()或named_buffers()看到。

4.2 遍历修改的几种姿势和坑

修改 buffer 的值，最常见方式是直接对张量做 inplace 操作，比如：

for buf in model.buffers(): buf.zero_()

因为 buffer 的requires_grad为 False，这里的 inplace 操作不会影响任何计算图。如果某个 buffer 的更新需要读取当前参数的值，也完全可行：

for name, buf in model.named_buffers(): if "running" in name: buf.copy_(some_other_tensor)

有的朋友习惯直接用buf = new_tensor这样重新赋值，但如果这个变量来自named_buffers()返回的元组，重新赋值通常只会改变局部变量，不会改写模块内部的_buffers。正确的修改姿势是使用copy_、fill_、mul_、add_这类 inplace 方法，或者用setattr(model_module, buffer_name, new_tensor)重新赋值。

如果 buffer 在 forward 里参与计算，更新它时才更要小心。比如在训练过程中对某些 buffer 做更新，但又不希望更新本身影响反向传播，这时就要明确地把更新逻辑包进torch.no_grad()：

def forward(self, x): if self.training: with torch.no_grad(): self.running_mean.lerp_(x.mean(dim=0), 0.1) return x

4.3 利用 buffers() 做模型状态对齐检查

我最近做了一件事：在不用跑推理的前提下，快速检查两个模型是否“状态一致”。光比较参数不够，buffer 也要比，因为 BN 统计量不同，推理结果就差很远。

下面的代码片段可以检查两个结构相同的模型在参数和 buffer 上是否完全一致：

def compare_models(model_a, model_b, device="cpu"): model_a = model_a.to(device).eval() model_b = model_b.to(device).eval() # 比较参数 for (name_a, pa), (name_b, pb) in zip( model_a.named_parameters(), model_b.named_parameters() ): if name_a != name_b: return False, f"参数名不一致: {name_a} vs {name_b}" if not torch.allclose(pa, pb): return False, f"参数值不一致: {name_a}" # 比较 buffer for (name_a, ba), (name_b, bb) in zip( model_a.named_buffers(), model_b.named_buffers() ): if name_a != name_b: return False, f"buffer名不一致: {name_a} vs {name_b}" if not torch.allclose(ba, bb): return False, f"buffer值不一致: {name_a}" return True, "参数和buffer完全一致"

这个函数在我排查问题的时候非常有用。比如某个模型在保存前和加载后推理结果不一致，直接跑一遍对比，十有八九是某个 buffer 的值变了。

5. 常见问题与排查技巧实录

把我在项目里见过的高频问题整理成速查表：

5.1 保存普通 Tensor 当成 buffer 使用，模型加载失效

这个坑我见过太多次了。有人不喜欢用register_buffer，直接在__init__里写：

self.running_sum = torch.zeros(...)

然后训练过程也正常更新它，model.to(device)时发现它没跟着走到 GPU，第一反应是“加个.cuda()”，于是运行时各种设备不一致。更隐蔽的是，保存模型时这个张量根本不进state_dict，加载后它永远是初始化值。排查方法非常简单：看看model.state_dict()里有没有这个 key，或者list(model.named_buffers())里有没有这个名字。如果没有，就是注册姿势不对。

5.2 persistent=False 的 buffer 在 load_state_dict 时表现

如果模型里注册了persistent=False的 buffer，那么state_dict()不会包含它，加载时也不会提示缺少它。这一点本身不是问题，问题在于如果你手动往 checkpoint 里塞了这个键，之后model.load_state_dict(checkpoint, strict=True)就会报unexpected key，因为模型当前state_dict里根本没有这个键。

我个人的建议是：除了临时调试用的状态外，尽量都用persistent=True，否则很容易忘掉“进程重启后这个 buffer 是全新初始状态”这件事，从而在长周期训练恢复时埋雷。

5.3 同名 buffer 重复注册会怎样

同一个模块里两次调用register_buffer注册同一个名字，后一次会覆盖前一次，不会报错。但如果这个名字和模块的某个已有属性冲突，就会报KeyError。

我在一个自动化代码生成项目里，出现过动态注册 buffer 时名字拼接错误、把两个不同层的统计量命名成同一个 key 的情况，结果后注册的覆盖了先注册的，模型内部数据直接错乱。这里给个硬性建议：动态注册 buffer 时，名字里所有路径信息都必须包含完整前缀，且不要用可变列表拼接，用固定格式生成。

5.4 把 buffer 塞进 optimizer 的错误操作

理论上，你可以把model.buffers()返回的张量拼成一个列表传给 optimizer。但这没有任何意义，因为 buffer 的requires_grad为 False，优化器内部会对这些参数做p.grad等操作，由于没有梯度，更新时要么跳过，要么在同一张量上反复执行zero_grad，行为非常混乱。

如果确实需要某个状态也参与优化，恰当的做法是把那个状态注册为nn.Parameter，而不是 buffer。语义上非常清晰：需要被梯度优化的，放参数里；不需要被梯度优化但需要带走的，放 buffer 里。

5.5 共享模块与递归遍历的去重行为

假设有一个子模块实例同时被两个父模块持有：

shared_bn = nn.BatchNorm2d(32) model = nn.Sequential( nn.Sequential(shared_bn), nn.Sequential(shared_bn), )

在较新版本 PyTorch 中，model.named_buffers()默认会去重，同一个 buffer 只返回一次。这个默认行为在大多数业务场景下没问题，但如果你想统计“每个持有者的 buffer 是否都一致”，想去重反而会漏掉信息。这时候可以通过remove_duplicate=False拿到每个引用位置各自看到的 buffer 对象。

这类共享逻辑通常在 NAS（神经架构搜索）或权重共享的层里出现。做这类模型的调试时，建议先把named_buffers(remove_duplicate=False)的结果打印出来，确认清楚每个名字对应哪个引用路径。

最后，聊聊我的一套自查习惯

接触 buffer 这个概念越深，就越觉得很多人其实只使用了nn.Module的表层功能。我后来在每次写完一个新的自定义层以后，都会强制自己走一遍这个自查流程：

• model.state_dict().keys()是否包含我预期的所有张量名称？
• model.to(device)之后，我手动print(model.buffers())检查设备是否一致？
• 训练完保存时，加载后跑一次 dummy forward，观察 BN 统计量有没有变化？
• 对比两个不同时间点的 checkpoint 时，到底 diff 的是参数还是 buffer？

每次走完这套流程，都要比单纯看堆栈报错省心得多。尤其是那个“普通 Tensor 属性 vs buffer”的坑，只要判断错一次，后面排查时间就没下限。记住那句最简单的话：需要在 checkpoint 里存活的状态张量，就老老实实register_buffer；能在模块间随着设备迁移的状态张量，也老老实实register_buffer。别图省事直接self.xxx = tensor，后面有的是更费劲的事在等你。