Rust in Everything
基于 burn 0.22.1

训练和推理模式

主线位置：训练主线里每个 step 之后的评估，和部署主线的推理：同一个模型切到推理行为，dropout 不丢、BatchNorm 用运行统计。
版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）
flowchart TB
  layer["层的 forward"] --> flag{"training 开着"}
  flag -->|"否"| inf["推理行为"]
  flag -->|"是"| dev{"输入在 autodiff 设备"}
  dev -->|"否"| inf
  dev -->|"是"| tr["训练行为"]
  ptflag["PyTorch 只看标志"] -.-> flag

实验代码：PyTorch 侧 dropout.py、bn.py、bn_sub.py、bn_clone.py，burn 侧 crate module-t 的 bin dropout、mixed、bnsub、bn、bnstale（burn =0.22.1，feature std、flex、ndarray、autodiff、optim，debug）。平台 M2 Max、macOS 27、rustc 1.99.0。

eval() 翻标志，valid() 换一个模型

PyTorch 的 Module.eval() 文档字符串说它「has an effect only on certain modules」，等价于 self.train(False)，返回 self：递归把每个子模块的 training 设成 False，原地生效。它不管求导，eval() 之后前向照样记图：

print("eval() returns self:", drop.eval() is drop)
model = a1_model().eval()
print("eval() model output requires_grad:", model(x).requires_grad)

eval() returns self: True
eval() model output requires_grad: True

所以 PyTorch 的推理要两件事一起做：model.eval() 加 torch.no_grad()。burn 0.22.1 把它们合成一个调用，而且返回新值：

PyTorchburn做了什么
model.eval()，原地let m = model.valid();，借用 &self 返回副本关掉 training 标志，参数去掉 autodiff
model.train()，原地let m = m.train();，按值拿走返回新模型恢复标志，参数抬回 autodiff、按配置要梯度
子模块 bn.eval()freeze() 再 set_require_grad(true)，见下文只关这个子模块的训练行为

Module::valid 的文档开头是「Returns a validation snapshot without autodiff or active training flags」。去掉 autodiff 那一半的效果（不记图、输出是普通张量）见不求导的区域；这一页讲标志那一半。Module::train 和 PyTorch 的 train() 也不完全一样：在普通 Device::flex() 上建的模型调 train()，参数会被抬到 autodiff 上并要梯度（实测见 burn 卷 Module 派生的「train、valid、freeze」）；PyTorch 的 train() 只翻标志。

因为 valid() 返回副本，训练循环里的写法是每次评估取一份，训练继续用原来的模型：

let valid_model = model.valid(); // model 还在，接着训练

Dropout：标志之外还看设备

PyTorch 的 nn.Dropout 只看 self.training。p=0.5，对 ones(1, 10000) 前向，数输出里的 0：

torch.manual_seed(0)
drop = nn.Dropout(0.5)
x = torch.ones(1, 10_000)

def show(label):
  out = drop(x)
  print(f"{label:34} zeros {int((out == 0).sum()):5}  max {out.max().item():g}  training={drop.training}")

show("fresh (train mode)")
drop.eval(); show("eval()")
drop.train()
with torch.no_grad():
  show("train(), inside no_grad")
with torch.inference_mode():
  show("train(), inside inference_mode")

fresh (train mode)                 zeros  5046  max 2  training=True
eval()                             zeros     0  max 1  training=False
train(), inside no_grad            zeros  5086  max 2  training=True
train(), inside inference_mode     zeros  4966  max 2  training=True

burn 侧同样的层，换几种状态和输入设备：

fn show(label: &str, drop: &Dropout, input_device: &Device) {
  let out = drop.forward(Tensor::<2>::ones([1, 10_000], input_device));
  let zeros: i64 = out.clone().equal_elem(0.0).int().sum().into_scalar();
  let max: f32 = out.max().into_scalar();
  println!("{label:40} zeros {zeros:5}  max {max}  training={}", drop.training.is_enabled());
}

let plain = Device::flex();
let ad = Device::flex().autodiff();
ad.seed(0);

let drop = DropoutConfig::new(0.5).init();
show("fresh, autodiff input", &drop, &ad);
show("fresh, plain input", &drop, &plain);
show("valid(), autodiff input", &drop.valid(), &ad);
show("valid(), plain input", &drop.valid(), &plain);
show("valid().train(), autodiff input", &drop.valid().train(), &ad);
show("no_grad(), autodiff input", &drop.clone().no_grad(), &ad);
show("freeze(), autodiff input", &drop.clone().freeze(), &ad);
show("train() on fresh, plain input", &drop.clone().train(), &plain);

fresh, autodiff input                    zeros  5048  max 2  training=true
fresh, plain input                       zeros     0  max 1  training=true
valid(), autodiff input                  zeros     0  max 1  training=false
valid(), plain input                     zeros     0  max 1  training=false
valid().train(), autodiff input          zeros  4969  max 2  training=true
no_grad(), autodiff input                zeros  4990  max 2  training=true
freeze(), autodiff input                 zeros     0  max 1  training=false
train() on fresh, plain input            zeros     0  max 1  training=true

两边丢弃的规则一样：丢掉约一半（5046、5048），留下的乘 1/(1-p) = 2。数值不同是随机数生成器不同。决定「丢不丢」的条件不一样，Dropout::forward 的第一行是：

        if !self.training.is_enabled() || !input.device().is_autodiff() || self.prob == 0.0 {
            return input;
        }
  • 标志：valid() 和 freeze() 关掉它，train() 恢复它，这部分和 PyTorch 的 eval() / train() 对应。no_grad() 不碰标志，所以照常丢弃，和 PyTorch「requires_grad_(False) 不影响 dropout」一致。
  • 输入设备：标志开着，输入不在 autodiff 设备上，也不丢（fresh, plain input）。在普通设备上做推理，忘了 valid() 也不会有 dropout；这在 PyTorch 里没有对应，PyTorch 的 no_grad()、inference_mode() 都不关 dropout，上面两行照样丢了一半。
  • 看的是到达 Dropout 的那个张量，不是批数据。模型在 autodiff 设备上、批数据建在 Device::flex() 上，第一层 Linear 的输出就已经在 autodiff 设备上了（参数在那里），后面的 Dropout 照常丢弃（bin mixed：Linear(4, 1000) → BatchNorm → Dropout(0.5)，数据在普通设备上时 8000 个输出里 3960 个是 0，BatchNorm 的运行均值也更新了）。只有直接作用在输入上的 Dropout 会因此失效。

对照起来，从 PyTorch 搬过来的几种写法：

PyTorchburndropout
model.train() 后前向模型和数据在 Device::flex().autodiff() 上，不调 valid()丢
model.eval() 后前向model.valid() 后前向，数据在哪都行不丢
model.train() + torch.no_grad()没有对应：普通设备上不丢，autodiff 设备上会记图PyTorch 丢
训练完直接在训练设备上评估、忘了 eval()同样会丢，见 burn 卷评估时结果每次不一样两边都丢
📝NOTE
冷知识：0.21 的 burn 真的没有训练/推理开关

0.21.0 的 Dropout::forward 只有一个条件：if !B::ad_enabled(&input.device()) || self.prob == 0.0 { return input; }，Dropout 结构体里只有 prob。那时「推理」就是把模型换到不带 autodiff 的后端上（valid() 返回内层后端上的另一个类型），没有可以单独翻的标志。PR #5498（「feat(module): generalize Param with module-owned flags」，首次发布于 v0.22.0-pre.4）给 Dropout、Noise、RRelu、BatchNorm 加了 training: Param<Flag>，条件变成上面那一行。新加的注释说明了为什么设备不够用：「The device says a backward is possible at all; the flag says this layer takes part in one, which a subtree frozen in place on the training device does not.」也就是部分微调时，冻结的子树还留在训练设备上，只看设备就分不出它该不该丢。所以「burn 没有 eval()，靠设备决定 dropout」这句话对 0.21 成立，对 0.22 只对一半。

只让一部分层用推理行为

PyTorch 微调时常见的写法：整个模型 train()，再对 BatchNorm 单独 eval()，让它用运行统计、不再更新，但 weight / bias 照样训练。burn 里 freeze 关掉标志的同时也关掉梯度，要再用 set_require_grad(true) 把梯度打开，它只改参数、不碰标志：

net = nn.Sequential(nn.Linear(4, 3), nn.BatchNorm1d(3))
net.train()
net[1].eval()  # 只让 BatchNorm 用运行统计，参数照样训练
x = torch.arange(32.0).reshape(8, 4) / 10
net(x).pow(2).sum().backward()
print("bn.weight.grad is None:", net[1].weight.grad is None, "| running_mean =", net[1].running_mean.tolist())
net.train()
print("after net.train(): bn.training =", net[1].training)

net.norm = match how.as_str() {
  "freeze" => net.norm.freeze(),
  "freeze+grad" => net.norm.freeze().set_require_grad(true),
  _ => unreachable!("freeze or freeze+grad"),
};
let x = Tensor::<1, burn::tensor::Int>::arange(0..32, &ad).float().reshape([8, 4]) / 10.0;
let y = net.norm.forward(net.linear.forward(x));
let grads = y.powi_scalar(2).sum().backward();
// 打印 gamma 有没有梯度、running_mean、training 标志
println!("{how:12} after net.valid().train(): training = {}", net.valid().train().norm.training.is_enabled());

bn.weight.grad is None: False | running_mean = [0.0, 0.0, 0.0]
after net.train(): bn.training = True

freeze       gamma grad is None: true | running_mean = [0.0, 0.0, 0.0] | training = false
freeze       after net.valid().train(): training = false
freeze+grad  gamma grad is None: false | running_mean = [0.0, 0.0, 0.0] | training = false
freeze+grad  after net.valid().train(): training = false
  • freeze().set_require_grad(true) 和 bn.eval() 效果相同：运行统计不动，gamma 有梯度。只写 freeze() 则 gamma 也不训练了。
  • 持久性相反。PyTorch 的 net.train() 递归地把 BatchNorm 也翻回训练模式，每个 epoch 开头调 model.train() 的循环要记得再 bn.eval() 一次。burn 的 freeze 改的是「配置上要不要训练」，train() 不会解开它，要恢复用 unfreeze()。

BatchNorm 的运行统计

PyTorch 的 BatchNorm1d(3)（momentum=0.1、eps=1e-5）和 burn 的 BatchNormConfig::new(3) 默认值相同，更新公式也相同：running = (1 - momentum) * running + momentum * batch_stat。同样三批数据（每批 4 个样本、3 个特征）在训练模式下前向，再切到推理，对比训练输出、第 1 批和第 3 批之后的运行统计、推理输出：

bn = nn.BatchNorm1d(3)
out["train_out"] = bn(batch(0)).tolist()
out["running_mean_1"] = bn.running_mean.tolist()
out["running_var_1"] = bn.running_var.tolist()
for k in (1, 2):
  bn(batch(k))
out["running_mean_3"] = bn.running_mean.tolist()
out["running_var_3"] = bn.running_var.tolist()
bn.eval()
out["eval_out"] = bn(batch(0)).tolist()

let bn = BatchNormConfig::new(3).init(&ad);
let train_out = rows(bn.forward(batch(0, &ad)));
let (m1, v1) = (mean(&bn), vec1(bn.running_var.value_sync()));
for k in 1..3 {
  bn.forward(batch(k, &ad));
}
let (m3, v3) = (mean(&bn), vec1(bn.running_var.value_sync()));
let eval_out = rows(bn.valid().forward(batch(0, &plain)));

train_out                n= 12  max|diff|=1.19e-07
running_mean_1           n=  3  max|diff|=0.00e+00
running_var_1            n=  3  max|diff|=1.24e-01
running_mean_3           n=  3  max|diff|=0.00e+00
running_var_3            n=  3  max|diff|=1.88e-01
eval_out                 n= 12  max|diff|=1.52e-01
worst=1.88e-01 tol=1e-06 MISMATCH

训练输出和运行均值对得上（1.19e-7 和 0，前者是归一化里 sqrt 和除法的舍入），运行方差差了 0.12 到 0.19，推理输出跟着差 0.15。第 1 批之后的运行方差：

PyTorch running_var_1 = [1.1664865, 1.3965679, 0.9717974]
burn    running_var_1 = [1.0998648, 1.2724260, 0.9538480]

原因在两边各自的文档里。PyTorch BatchNorm1d 的文档字符串：「At train time in the forward pass, the variance is calculated via the biased estimator, equivalent to torch.var(input, correction=0). However, the value stored in the moving average of the variance is calculated via the unbiased estimator, equivalent to torch.var(input, correction=1).」burn 的 ModuleOps::batch_norm_train 的文档说它返回「the returned mean and biased variance — what a training path feeds its running statistics」，BatchNorm::forward_train 把这个有偏方差直接混进运行方差：

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

归一化两边都用有偏方差，所以训练输出一致；存进运行方差的，PyTorch 是无偏的（除以 n-1），burn 是有偏的（除以 n）。每批 4 个样本，批方差差 4/3 倍：0.9 + 0.1 × 2.6649 = 1.1665，0.9 + 0.1 × 2.6649 × 3/4 = 1.0999。把 PyTorch 侧的运行方差改成用 var(unbiased=False) 手算一遍，和 burn 的输出再对比：

train_out                n= 12  max|diff|=1.19e-07
running_mean_1           n=  3  max|diff|=0.00e+00
running_var_1            n=  3  max|diff|=0.00e+00
running_mean_3           n=  3  max|diff|=0.00e+00
running_var_3            n=  3  max|diff|=0.00e+00
eval_out                 n= 12  max|diff|=0.00e+00
worst=1.19e-07 tol=1e-06 OK

全部对上，差别只在这一处。影响有多大取决于每个通道的样本数 n（批大小乘空间尺寸）：批方差项差 n/(n-1) 倍，n 上千时在千分之一以下，全连接层配小批量时最明显。从 PyTorch 导入训练好的运行统计不受影响（存的就是那个数）；在 burn 里接着训练或从头训练，运行方差会系统地偏小，推理输出和 PyTorch 训练出来的不一样。origin/main 上这两处没有改。

哪些情形会更新运行统计（bin bn 的后半段，每种一个新层、前向一次第 0 批）：

autodiff init, autodiff input    running_mean = [0.009881697, 0.09698707, 0.1868625]
autodiff init, plain input       running_mean = [0.0, 0.0, 0.0]
plain init, plain input          running_mean = [0.0, 0.0, 0.0]
no_grad(), autodiff input        running_mean = [0.009881697, 0.09698707, 0.1868625]
freeze(), autodiff input         running_mean = [0.0, 0.0, 0.0]

train() + no_grad: running_mean = [0.009882, 0.096987, 0.186862]

BatchNorm::forward 用和 Dropout 一样的两个条件选路：输入在 autodiff 设备上且标志开着，才走训练路径并更新统计。PyTorch 侧最后一行是训练模式下套 torch.no_grad()：照样更新，和 burn 在普通设备上「不更新」相反。

不调 valid() 的推理路径会读到旧统计

训练路径的更新经过 RunningState::update：新值先按线程放进一张表，下一次 update 或 value_sync 时才合并（多线程时取平均）。推理路径 forward_inference 读的是 value，文档写着「The current value might be outdated by one update」。valid() 会先同步，所以走 valid() 没有问题；不调 valid()、只把输入换成普通设备，就会少最近一次更新（bin bnstale，训练前向一次后，用同一层推理第 0 批，打印第一行输出）：

let bn = BatchNormConfig::new(3).init(&ad);
bn.forward(batch(&ad)); // one training step: updates the running stats
println!("running_mean.value()      = {:?}", vec1(bn.running_mean.value()));
println!("plain input, no valid()   = {:?}", rows(bn.forward(batch(&plain)))[0]);
println!("running_mean.value_sync() = {:?}", vec1(bn.running_mean.value_sync()));
println!("plain input, after sync   = {:?}", rows(bn.forward(batch(&plain)))[0]);
println!("valid(), plain input      = {:?}", rows(bn.valid().forward(batch(&plain)))[0]);

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]

第二行用的还是初始的均值 0 和方差 1；第四行只同步了均值，方差还是旧的；只有 valid() 两个都同步了。结论和 Dropout 一样：评估和推理一律先 valid()，不要靠「输入放在普通设备上」切换行为。

❗IMPORTANT
动脑筋：clone 出来的模型前向一次，原件的运行均值变不变？

PyTorch 里 copy.deepcopy(bn)(x) 不影响 bn。burn 里：

let bn = BatchNormConfig::new(3).init(&ad);
let copy = bn.clone();
copy.forward(batch(0, &ad));
println!("forward on a clone: original running_mean = {:?}", mean(&bn));

打印的是 [0.0, 0.0, 0.0] 吗？

动脑筋答案

forward on deepcopy: original running_mean = [0.0, 0.0, 0.0]
forward on the same object: running_mean = [0.009882, 0.096987, 0.186862]

forward on a clone: original running_mean = [0.009881697, 0.09698707, 0.1868625]
  • 变了，和对原件本身前向一次的结果一样。RunningState 里是 Arc<Mutex<...>>，派生出来的 Clone 只复制 Arc，克隆和原件共用同一份运行统计。这是有意的：BatchNorm::forward 只拿 &self，多线程各自前向时更新要汇到同一处，update 再按线程取平均。
  • 参数没有这个问题：Param 的值不可变，优化器 step 返回的是新模型。会被前向改掉的只有运行统计。
  • 想要一份独立的 BatchNorm，用 valid()（它新建 RunningState）或者重新 init 再加载记录；不要指望 clone() 像 deepcopy。

小结

  • model.eval() 原地翻标志、不管求导；burn 的 valid() 返回副本，同时关标志、去 autodiff；train() 按值拿走模型，普通设备上的模型也会被抬到 autodiff。
  • Dropout 和 BatchNorm 的训练行为要「标志开着」且「到达该层的张量在 autodiff 设备上」；PyTorch 只看标志，no_grad() 不关 dropout。
  • 子模块 bn.eval() 对 freeze().set_require_grad(true)；burn 的冻结跨过 train() 仍有效，PyTorch 的 train() 会把它翻回去。
  • BatchNorm 训练输出一致（1.19e-7）；运行方差 burn 用有偏估计（除以 n），PyTorch 用无偏估计（除以 n-1），4 个样本一批、三批之后推理输出差 0.15。
  • 推理先 valid()：不调它时 BatchNorm 读到的运行统计少最近一次更新；clone() 共用运行统计。

💬 关联讨论

发起讨论
加载中...