验证循环里内存一直涨
主线位置：训练主线的旁路：每个 epoch 训练完跑一遍验证集，只有 forward，没有 backward。部署主线的「推理」一步也是同一回事。
版本：PyTorch 2.15.0（CPU，
torch.set_num_threads(1)）、burn 0.22.1（Flex 设备加 autodiff，release 构建）
sequenceDiagram participant T as 训练一步 participant V as 验证 50 批 participant G as 计算图 T->>G: backward 清理 V->>G: 每批 forward 留下节点 Note over G: 涨约 1 MiB 每批 T->>G: 下一次 backward 才清理 V->>G: 训练结束后的验证 Note over G: 再也没有 backward
实验 crate pita-t 的 p11，PyTorch 侧 p11_eval.py。模型是 Linear(512, 512)、ReLU、Linear(512, 10)；训练 batch [64, 512]，验证 batch [256, 512]；3 个 epoch，每个 epoch 训练一步、验证 50 批，最后再验证 200 批。RSS 用 ps -o rss= 读，相对开始时的增量。两边都固定了种子（burn ad.seed(42)，PyTorch torch.manual_seed(42)），但两边的随机数不通用，只看各自的变化。平台 M2 Max、macOS 27、rustc 1.99.0。
PyTorch 的写法
正规写法是 model.eval() 加 torch.no_grad()。忘了 no_grad 也能跑：
def evaluate(n):
score = 0.0
for _ in range(n):
if mode == "no_grad":
with torch.no_grad():
out = model(val)
else:
out = model(val)
score += out.mean().item()
return score
plain epoch 1 after train step: rss + 6.5 MiB
plain epoch 1 after 50 val batches: rss + 11.1 MiB
plain epoch 2 after train step: rss + 11.1 MiB
plain epoch 2 after 50 val batches: rss + 11.1 MiB
plain epoch 3 after train step: rss + 11.1 MiB
plain epoch 3 after 50 val batches: rss + 11.1 MiB
plain final eval, 200 batches: rss + 11.1 MiB
no_grad epoch 1 after train step: rss + 7.5 MiB
...
no_grad final eval, 200 batches: rss + 10.7 MiB
不加 no_grad 时每批会建图，但 out 在下一轮被覆盖，图挂在 out.grad_fn 上，跟着一起释放。内存是平的，代价只是多算了一点。
搬到 burn 的现象
模型建在 Device::flex().autodiff() 上用来训练，验证时直接拿同一个模型 forward：
for epoch in 1..=3 {
let loss = model.forward(x.clone()).powi_scalar(2).mean();
let grads = GradientsParams::from_grads(loss.backward(), &model);
model = optim.step(1e-3, model, grads);
let mut score = 0.0f32;
for batch in 0..50 {
let out = model.forward(val.clone());
score += out.mean().into_scalar::<f32>();
}
}
// 训练结束后的最终验证：200 批，之后没有 backward
autodiff epoch 1 after train step: rss + 5.2 MiB
autodiff val output: is_autodiff=true is_tracked=true
autodiff epoch 1 after 50 val batches: rss + 55.5 MiB (score 0.341)
autodiff epoch 2 after train step: rss + 57.2 MiB
autodiff epoch 2 after 50 val batches: rss + 65.3 MiB (score 0.338)
autodiff epoch 3 after train step: rss + 65.3 MiB
autodiff epoch 3 after 50 val batches: rss + 66.8 MiB (score 0.335)
autodiff final eval, 200 batches: rss + 258.1 MiB
- 第一轮 50 批验证，RSS 涨了约 50 MiB，每批约 1 MiB，相当于两个
[256, 512]的 f32 张量（各 0.5 MiB）；前向里fc1的输出和 ReLU 的输出就是这个形状。 - 第二、三轮不再明显上涨：中间那次训练的 backward 清掉了上一轮验证留下的图，内存被复用（RSS 本身不降，分配器不把页还给系统）。
- 训练结束后的 200 批验证没有后续的 backward，一路涨到 +258 MiB。验证集更大、模型更大时，这里就是 OOM。
输出的 is_tracked=true 说明验证的前向在记图：模型参数要梯度，图就从参数长出来。
原因
PyTorch 的图挂在输出张量上，输出没人引用就释放。burn 的图存在进程级的结构里，AutodiffTensor 没有 Drop 去通知它，节点只在 backward 结束时被清理。所以在 autodiff 设备上只做前向，图一直涨到下一次 backward；一直没有 backward，就一直涨。机制和源码见不求导的区域的最后一节和 burn 卷推理时内存一直涨，那一页还量了「完全没有 require_grad 时每个算子约 1 KB」的情形。
burn 也没有一个像 torch.no_grad() 那样的作用域可以包住验证循环：记不记图看的是张量所在的设备，不看调用的位置。
改法
验证前用 Module::valid() 拿一份不带 autodiff 的模型，输入也去掉 autodiff：
let out = model.valid().forward(val.clone().without_autodiff());
valid epoch 1 after train step: rss + 5.1 MiB
valid val output: is_autodiff=false is_tracked=false
valid epoch 1 after 50 val batches: rss + 6.8 MiB (score 0.341)
valid epoch 2 after train step: rss + 7.7 MiB
valid epoch 2 after 50 val batches: rss + 7.7 MiB (score 0.338)
valid epoch 3 after train step: rss + 7.7 MiB
valid epoch 3 after 50 val batches: rss + 7.7 MiB (score 0.335)
valid final eval, 200 batches: rss + 7.7 MiB
内存平了，三轮的 score 和 autodiff 版本逐位相同：valid() 只是去掉 autodiff（顺带关掉 dropout 这类训练行为，见训练和推理模式），前向的数值不变。valid() 借用 &self 返回一个新模型。实验里每批调一次，内存和数值都没有问题；放在验证循环外面调一次更清楚。
部署时更直接：模型一开始就建在不带 autodiff 的设备上，或者加载后 valid()，见部署。
怎么检查
- 验证、推理入口处断言
!out.is_tracked()，或者!x.device().is_autodiff()。 - 长时间的验证或推理，看 RSS 每批涨不涨。burn 的「涨到下一次 backward」在训练循环里会被下一步的 backward 掩盖，最容易在训练结束后的最终评估、生成样本、导出前的冒烟测试里暴露。
- 训练和推理别共用一个
device变量：Device的相等比较不看 autodiff（burn 卷推理时内存一直涨的修法一节），推理代码单独建一个Device::flex()。
小结
- PyTorch 忘了
no_grad的验证不会涨内存，图随输出释放；burn 在 autodiff 模型上 forward，图留到下一次 backward。 - 实测每个验证 batch 约 1 MiB；训练结束后 200 批涨到 +258 MiB。
- 验证用
model.valid()加without_autodiff()的输入，数值不变，内存平。