BatchNorm 推理结果和 PyTorch 对不上

主线位置：训练主线每一步的 forward 顺带更新运行统计，训练结束后的评估和部署主线的推理读这些统计。

版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备，autodiff）

flowchart TB
  tr["训练 forward"] --> upd["更新 running_var"]
  upd -->|"有偏 除以 n"| rv["运行方差偏小"]
  rv --> ev["推理输出差"]
  snap["clone 快照"] -.->|"共用 Arc"| upd
  ev2["不调 valid 推理"] -.->|"少一次合并"| upd

实验代码：PyTorch 侧 bnpit.py、bn_eval1.py，burn 侧 crate pitb-module 的 bin bnpit、bnstale（burn =0.22.1，feature std、flex、ndarray、autodiff、optim，debug）。一个 BatchNorm1d(3) / BatchNormConfig::new(3)（两边默认 momentum 0.1、eps 1e-5），每批 n 个样本，数据是 2·sin(3r + c + 1 + k) + c（f64 里算、舍入成 f32 一次，两边逐位相同）。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch 的写法

训练 20 批；训练到一半用 copy.deepcopy 留一份快照（「保存目前最好的模型」的常见写法）；最后 eval()，推理一批：

bn = nn.BatchNorm1d(3)
for k in range(10):
  bn(batch(k))
snap = copy.deepcopy(bn)
snap_mean = snap.running_mean.tolist()
for k in range(10, 20):
  bn(batch(k))
print(f"n={n} snapshot running_mean unchanged: {snap.running_mean.tolist() == snap_mean}")
bn.eval(); snap.eval()
x = batch(100)[:4]
out = {
  "running_mean": bn.running_mean.tolist(), "running_var": bn.running_var.tolist(),
  "snap_running_mean": snap.running_mean.tolist(),
  "eval_out": bn(x).tolist(),
}

n=4 snapshot running_mean unchanged: True
n=4 running_var = [2.427054, 2.61856, 2.33988]
n=64 snapshot running_mean unchanged: True
n=64 running_var = [1.90205, 1.903964, 1.91132]

搬到 burn 的现象

同样的 20 批，快照用 clone()，推理前 valid()（bin bnpit）：

let mut bn = BatchNormConfig::new(3).init(&ad);
for k in 0..10 {
  bn.forward(batch(k, n, &ad));
}
let snap: BatchNorm = match how.as_str() {
  "clone" => bn.clone(),
  "valid" => bn.valid(),
  // ...
};
let snap_mean = vec1(snap.running_mean.value_sync());
for k in 10..20 {
  bn.forward(batch(k, n, &ad));
}
let snap_after = vec1(snap.running_mean.value_sync());
println!("n={n} {how}: snapshot running_mean unchanged: {}", snap_mean == snap_after);
let x = batch(100, 4, &plain);
// 输出 running_mean、running_var、快照的 running_mean、bn.valid().forward(x)

n=4 clone: snapshot running_mean unchanged: false
n=4 running_var = [1.850684, 1.994314, 1.785304]

和 PyTorch 的输出逐项比（n = 4）：

running_mean             n=  3  max|diff|=9.31e-10
running_var              n=  3  max|diff|=6.24e-01
snap_running_mean        n=  3  max|diff|=4.54e-01
eval_out                 n= 12  max|diff|=1.90e-01
worst=6.24e-01 tol=1e-06 MISMATCH

两件事：

• 运行均值对得上，运行方差差 0.62，推理输出跟着差 0.19。 每批样本越多差得越少，n = 64 时运行方差差 2.80e-2、推理输出差 1.13e-2，还是远超 1e-6。
• clone() 出来的快照不是快照。 原模型后来又训练了 10 批，快照的运行均值跟着变了，和 PyTorch deepcopy 留下的值差 0.45。

还有第三种，推理时不调 valid()，只把输入放在普通设备上（bin bnstale：训练前向一次，然后用同一层推理第一行）：

running_mean.value()      = [0.0, 0.0, 0.0]
plain input, no valid()   = [1.6829334, 2.8185809, 2.2822285]
running_mean.value_sync() = [0.009881697, 0.09698707, 0.1868625]
plain input, after sync   = [1.6730518, 2.7215943, 2.095367]
valid(), plain input      = [1.595291, 2.4127238, 2.14546]

PyTorch 同样训练一次后 eval()，推理同一行：

eval() first row = [1.5490658283233643, 2.302994728088379, 2.125554323196411]

不调 valid() 的那一行用的是初始统计（均值 0、方差 1）；valid() 那一行剩下的差就是上面第一件事。

原因

三件事的机制在训练和推理模式的「BatchNorm 的运行统计」里都实测过，这里只列出处：

• 有偏方差。 BatchNorm::forward_train 把 batch_norm_train 返回的批方差直接混进运行方差：

  Copy

          let running_var = running_var
              .mul_scalar(1.0 - self.momentum)
              .add(result.variance.detach().mul_scalar(self.momentum));
  

  这个 variance 是有偏的（除以 n），PyTorch 存进 running_var 的是无偏的（除以 n - 1，BatchNorm1d 的文档字符串写明了「the value stored in the moving average of the variance is calculated via the unbiased estimator」）。归一化本身两边都用有偏方差，所以训练输出和运行均值对得上，只有运行方差每一步少乘一个 n/(n-1)。BatchNormConfig 没有切换的选项。

• clone() 共用统计。 RunningState 里是 Arc<Mutex<..>>，派生的 Clone 只复制 Arc。forward 只拿 &self，统计靠这把锁改，克隆和原件改的是同一份。参数不受影响：参数值不可变，优化器 step 返回新模型。

• 少一次更新。 update 先把新值按线程存起来，下一次 update 或 value_sync 才合并；推理路径 forward_inference 读的 value 不合并，文档写着「The current value might be outdated by one update」。RunningState 的 valid 会先合并，再新建一个 RunningState。

origin/main 上 BatchNorm 只是把选路条件提成了一个 trains 方法、加了性能统计的标注，有偏方差、RunningState 的 Clone 和 value 都没有改。有偏方差是 burn 0.22.1 和 PyTorch 的行为差异，burn 的文档（batch_norm_train 的「biased variance — what a training path feeds its running statistics」）说明它是有意的。

改法

快照和推理都用 valid()。 它新建 RunningState，和原模型从此无关；推理前调它，统计是合并过的：

n=4 valid: snapshot running_mean unchanged: true

snap_running_mean        n=  3  max|diff|=0.00e+00

valid() 的快照不在 autodiff 设备上，要接着训练它就再 .train()。

要和 PyTorch 训练出的推理结果对齐，换算运行方差。 每批样本数固定为 n、共训练 K 步时，运行方差里来自初始值 1 的部分是 c = (1 - momentum)^K，其余部分每一项都差 n/(n-1) 倍，可以一次换算回来：

let c = (1.0 - bn.momentum).powi(20);
let corrected = (bn.running_var.value_sync() - c) * (n as f64 / (n as f64 - 1.0)) + c;
bn.running_var = RunningState::new(corrected);

n=4 fix: snapshot running_mean unchanged: true
n=4 running_var = [2.427054, 2.61856, 2.33988]
running_mean             n=  3  max|diff|=9.31e-10
running_var              n=  3  max|diff|=2.38e-07
snap_running_mean        n=  3  max|diff|=0.00e+00
eval_out                 n= 12  max|diff|=1.19e-07
worst=2.38e-07 tol=1e-06 OK

n = 64 换算后推理输出最大差 3.58e-7，同样通过。这个换算只在每批样本数都相同时成立（最后一批不满就不对了），它证明了差别全在这一处，不是通用的修法。实际取舍：

• 从 PyTorch 导入训练好的 BatchNorm（见导入 PyTorch 权重）不受影响，文件里存的就是 PyTorch 的无偏值。
• 在 burn 里训练、在 burn 里推理，有偏的运行方差是自洽的，只是和同样训练的 PyTorch 模型不一致；每个通道的样本数（批大小乘空间尺寸）上千时差别在千分之一以下。
• 要求两边结果一致的场景（这一卷的对照、迁移验证），按上面换算，或在 PyTorch 侧对比时改用有偏方差手算（训练和推理模式里这样对上了 0 误差）。

怎么检查

• 把 copy.deepcopy(model) 换成 model.clone() 的地方逐个看：模型里有 BatchNorm，就改成 model.valid()。
• 推理和评估代码里找没有经过 valid() 的 forward；「输入放在普通设备上」不等于推理模式。
• 迁移时对比运行统计：running_mean 对上而 running_var 偏小、比值接近 (n-1)/n，就是有偏方差。
• 模型里有 BatchNorm 时，对比 num_params() 记得它多算了运行统计（见参数与遍历）。

小结

• burn 的 BatchNorm 运行方差用有偏估计：同样 20 批、每批 4 个样本，推理输出和 PyTorch 差 0.19；每批 64 个样本差 1.13e-2；运行均值对得上。
• clone() 和原件共用运行统计，快照要用 valid()；deepcopy 的对应不是 clone()。
• 推理前不调 valid()，BatchNorm 读到少一次合并的统计。
• 批大小固定时把运行方差按 n/(n-1) 换算，推理输出和 PyTorch 最大差 1.19e-7（n = 4）。