📚 课程目录第3章 BN的训练与推理模式

3. BN 的训练与推理模式：moving mean 与 moving variance

说白了，Batch Normalization 在训练和推理时，是两套完全不同的逻辑。很多人一开始只记得前向传播那一套——减均值、除方差、再缩放平移，结果模型部署上线后，推理结果和训练时对不上，一脸懵。这中间的坑，就在 moving mean 和 moving variance 这两个参数上。

我个人习惯把 BN 的训练模式叫做“实时统计模式”，推理模式叫做“冻结统计模式”。咱们今天就把它彻底拆开来看。

3.1 训练阶段：实时统计与滑动更新

训练时，每一层的 BN 都会做两件事：

当前 batch 的归一化：用当前 batch 的 mean 和 variance 做归一化。
全局统计量的滑动更新：用当前 batch 的统计量去更新 moving mean 和 moving variance。

注意，第二个动作不参与梯度反传，它只是“顺便”记录下来。更新的公式很简单：

moving_mean = momentum * moving_mean + (1 - momentum) * batch_mean
moving_var  = momentum * moving_var  + (1 - momentum) * batch_var

这里的 momentum 默认通常是 0.9 或 0.99。值越大，历史统计量保留得越多，更新越平滑。值越小，模型对近期 batch 的变化越敏感。

我的小习惯：如果在分子生成任务中看到 loss 曲线剧烈震荡，可以检查一下 BN 的 momentum 是不是设得太小了。我遇到过一种情况——某个分子数据集样本不平衡，一个 batch 里全是活性分子，下一个 batch 全是非活性分子，这时 momentum=0.9 就不够稳定，我后来调到 0.95，训练就稳下来了。

为什么要同时维护这两套统计量？你想想看，训练时 batch size 一般比较大（比如 32 或 64），batch 内的统计量基本能代表整体分布。但到了推理时，如果还依赖 batch 内的统计量，就出大问题了——推理时通常一次只输入一个样本，batch size = 1，那 mean 就是这个样本本身，variance 直接是 0，归一化就毫无意义了。

3.2 推理阶段：冻结全局统计量

推理时，BN 的行为很简单：

不再计算当前 batch 的 mean 和 variance。
直接使用训练阶段累积的 moving mean 和 moving variance。
缩放和平移参数（γ 和 β）依然是训练学到的值。

PyTorch 里调用 model.eval() 时，BN 层就会自动切换到推理模式。如果你忘了切，或者代码里在推理时没调 model.eval()，BN 依然会按 batch 统计量归一化，结果就是每个样本的归一化结果完全依赖它所在的 batch 构成，输出会随机抖动。这在分子生成任务里是致命的——同样的分子输入，两次推理可能得到不同的特征表示。

我曾经踩过的坑：有一次做分子性质预测，模型在验证集上跑了两次，结果不一致。排查了半天，原来是在推理脚本里忘了调 model.eval()，BN 一直在用 batch 统计量。特别是 batch size = 1 时，variance = 0，除零操作直接导致输出全变成 NaN。花了整整一个下午才定位到这个问题。

3.3 moving mean 与 moving variance 的初始化与保存

这两个参数的初始值是多少？

moving_mean 初始化为 0。
moving_variance 初始化为 1。

为什么？因为在没有看到任何数据之前，我们假设输入是标准正态分布。随着训练的进行，这两个值会逐渐学习到真实的全局分布。

模型保存时，这两个参数和 γ、β 一样，都是 state_dict 的一部分。很多同学只保存了模型参数，却没意识到 moving mean 和 moving variance 也是“参数”——它们是训练过程中学到的，不是推理时计算的。

# PyTorch 中查看 BN 层的参数
for name, param in bn_layer.named_parameters():
    print(name, param.shape)

# 输出：
# weight (gamma)     torch.Size([256])
# bias (beta)        torch.Size([256])

# 但 moving_mean 和 moving_var 不是 parameters，是 buffers
for name, buffer in bn_layer.named_buffers():
    print(name, buffer.shape)

# 输出：
# running_mean       torch.Size([256])
# running_var        torch.Size([256])
# num_batches_tracked torch.Size([])  # 可选

这里有个细节：num_batches_tracked 记录的是看到过的 batch 数，某些框架用这个来做无偏估计。不过大多数情况下，直接用 moving mean 和 moving variance 就够了。

  关键区别总结：BN 的 γ 和 β 是可学习参数，通过梯度下降优化；而 moving mean 和 moving variance 是统计量，通过滑动平均累积。前者是“学”出来的，后者是“算”出来的。

3.4 显式对比：训练 vs 推理

特性	训练模式	推理模式
计算 mean	当前 batch 实时计算	使用 running_mean
计算 variance	当前 batch 实时计算	使用 running_var
更新 running 参数	滑动更新	不做更新
γ、β	参与梯度更新	冻结
batch size 敏感	对 batch size 敏感，太小时不稳定	完全不依赖 batch size
分子生成场景	大 batch 训练，分布稳定	单个分子推理，必须用全局统计量

3.5 为什么不直接用整个训练集的 mean/var？

有人可能会问：既然推理时要用全局统计量，为什么不提前算好整个训练集的均值和方差，直接存起来？

原因有几点：

计算代价高：分子数据量通常很大，几百万个分子过一遍前向，只为算 mean 和 var，太浪费。
分布漂移：训练过程中，前一层的输出分布会随着参数更新而变化。用固定的全数据集统计量，跟不上这种变化。
滑动平均更平滑：滑动平均天然对近期 batch 有更高的权重，能更好地反应模型当前状态下的数据分布。

补充一个经验：在分子生成模型中，如果训练集和测试集的分子结构差异很大（比如训练集全是小分子，测试集大分子多），推理时 BN 的效果可能会变差。这时候可以尝试在推理时做一次“统计量校准”——用一个小的验证集重新计算一次 running_mean 和 running_var，再冻结住。我在实际项目中用过这个技巧，效果还不错。

3.6 框架中的实现差异

不同深度学习框架对 BN 推理模式的实现略有差异，但核心逻辑一样：

# PyTorch：通过 train/eval 模式切换
model.train()   # BN 使用 batch 统计量 + 更新 running 参数
model.eval()    # BN 使用 running 统计量，不更新

# TensorFlow/Keras：通过 training 参数控制
# training=True：batch 统计量
# training=False：running 统计量（默认）
def call(self, x, training=False):
    if training:
        # 计算 batch mean/var，更新 running
    else:
        # 直接使用 running

# JAX/Flax：通常需要显式传递 use_running_average=True
# 或通过 mutable 变量控制

嗯，这里要注意：PyTorch 里 model.eval() 会让所有 BN 层切换模式，但也会同时让 Dropout 失效。如果你只想切换 BN 而不影响 Dropout，可以用 torch.no_grad() 配合 bn_layer.training = True 之类的骚操作，不过我不建议这么搞，容易搞混。代码可读性更重要。

3.7 一张图理解整个过程

下面的 SVG 展示了 BN 在训练和推理两个阶段的完整流程对比，核心就是统计量的来源不同：

回到最初的问题。BN 的训练和推理模式其实就一句话：训练时即算即用、边走边记；推理时直接翻旧账、不写新账。理解了这个本质，就永远不会在部署时被 BN 坑到。