训练循环

主线位置：训练主线的最外层：把一步放进 epoch 循环，每个 epoch 末验证、记指标、存 checkpoint。

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

flowchart TB
  subgraph hand ["手写循环 两边都能写"]
    h1["for epoch"] --> h2["for batch 一步"]
    h2 --> h3["切到验证模式 算指标"]
    h3 --> h4["存 checkpoint"]
  end
  subgraph lr ["burn Learner"]
    l1["TrainStep step"] --> l2["SupervisedTraining"]
    l2 --> l3["指标线程 写日志"]
    l2 --> l4["三个 checkpoint 文件"]
    l2 --> l5["LearningResult"]
  end
  h2 -.->|"同一串调用"| l1

实验代码：PyTorch 侧 loop.py，burn 侧 crate train-loop（模板的 feature 加 train，train 会带上 dataset；default-features = false，所以没有 TUI），对比脚本 check.py。release 构建。平台 M2 Max、macOS 27、rustc 1.99.0。

同一个任务

模型和初值是 A1 的两层 MLP。数据是合成的三分类：第 k 个样本的四个特征是 sin(0.37k + 1.3j)（j = 0..3），先在 f64 里算再转 f32，所以两边逐位相同；标签取 [x0 + x1, x2 − x3, x3 − x0] 里最大的那个。前 96 个训练、后 24 个验证，每批 16 个、不打乱，SGD 学习率 0.1，3 个 epoch，共 18 步。验证集最后一批只有 8 个。

PyTorch 的手写循环：

opt = torch.optim.SGD(model.parameters(), lr=0.1)
for epoch in range(1, EPOCHS + 1):
  model.train()
  for i in range(0, len(xt), BATCH):
    loss = nn.functional.cross_entropy(model(xt[i:i + BATCH]), yt[i:i + BATCH])
    losses.append(loss.item())
    opt.zero_grad()
    loss.backward()
    opt.step()
  model.eval()
  with torch.no_grad():
    s = ws = n = correct = batches = 0
    for i in range(0, len(xv), BATCH):
      out = model(xv[i:i + BATCH])
      l = nn.functional.cross_entropy(out, yv[i:i + BATCH]).item()
      m = len(yv[i:i + BATCH])
      s, ws, n, batches = s + l, ws + l * m, n + m, batches + 1
      correct += (out.argmax(1) == yv[i:i + BATCH]).sum().item()
  v_mean.append(s / batches)
  v_weighted.append(ws / n)
  v_acc.append(100.0 * correct / n)
  torch.save({"epoch": epoch, "model": model.state_dict(), "optim": opt.state_dict()}, f"out/ckpt/epoch-{epoch}.pt")

验证 loss 记了两种平均：v_mean 是各批 loss 的算术平均，v_weighted 按每批样本数加权。后面会用到。

burn 的手写循环几乎一行对一行。数据用 burn 的 DataLoaderBuilder（Batcher 和加载器的写法见 Dataset 与 DataLoader）：

let train = DataLoaderBuilder::new(SampleBatcher)
  .batch_size(BATCH)
  .set_device(device.clone())
  .build(InMemDataset::new(samples(0..96)));
let valid = DataLoaderBuilder::new(SampleBatcher)
  .batch_size(BATCH)
  .set_device(device.clone().inner())
  .build(InMemDataset::new(samples(96..120)));
let mut model = model0.clone();
let mut optim = SgdConfig::new().init();
for _epoch in 1..=EPOCHS {
  for batch in train.iter() {
    let out = model.forward_classification(batch.expect("in-memory dataset"));
    losses.push(out.loss.clone().into_scalar::<f32>());
    let grads = GradientsParams::from_grads(out.loss.backward(), &model);
    model = optim.step(LR, model, grads);
  }
  let eval = model.valid();
  for batch in valid.iter() {
    let batch = batch.expect("in-memory dataset");
    let len = batch.y.dims()[0];
    let out = eval.forward_classification(batch);
    // ... 同样累加 sum、wsum、n、correct、batches
  }
  // ... 同样算 v_mean、v_weighted、v_acc
}

forward_classification 是模型上的一个方法，前向加交叉熵，返回 burn 的 ClassificationOutput（loss、output、targets 三个张量）。两处和 PyTorch 不同：

没有 model.train()，验证前是 let eval = model.valid();：得到一份去掉 autodiff 的副本，在它上面前向不记图，model 本身不变，下一个 epoch 接着训练。这对应 PyTorch 的 model.eval() 加 torch.no_grad() 两件事（见不求导的区域、训练和推理模式）。
验证集的加载器建在 device.inner() 上，也就是不带 autodiff 的同一个设备。

Learner 版本：模型实现 TrainStep 和 InferenceStep，其余交给 SupervisedTraining：

impl TrainStep for Mlp {
  type Input = Batch;
  type Output = ClassificationOutput;

  fn step(&self, batch: Batch) -> TrainOutput<ClassificationOutput> {
    let out = self.forward_classification(batch);
    if let Ok(mut v) = LEARNER_LOSSES.lock() {
      v.push(out.loss.clone().into_scalar());
    }
    TrainOutput::new(self, out.loss.backward(), out)
  }
}

impl InferenceStep for Mlp {
  type Input = Batch;
  type Output = ClassificationOutput;

  fn step(&self, batch: Batch) -> ClassificationOutput {
    self.forward_classification(batch)
  }
}

let training = SupervisedTraining::new(dir, train, valid)
  .metric_train_numeric(LossMetric::new())
  .metric_valid_numeric(LossMetric::new())
  .metric_valid_numeric(AccuracyMetric::new())
  .with_default_checkpointers()
  .num_epochs(EPOCHS)
  .renderer(Box::new(CliMetricsRenderer::new()))
  .summary();
let result = training.launch(Learner::new(model0.clone(), SgdConfig::new().init(), LR));

LEARNER_LOSSES 是一个全局的 Mutex<Vec<f32>>，只为了把 Learner 每一步的 loss 拿出来对比，正常使用不需要。这里的两个加载器没有 set_device：Learner 的单设备策略会把训练加载器搬到 autodiff 设备、验证加载器搬到 device.inner()。TrainOutput::new 里就是 GradientsParams::from_grads，每一步的拆解见 burn 卷的 Learner 与数据。

三份 loss

check.py 先用 compare.py 比 PyTorch 和 burn 的手写循环，再逐位（struct.pack("<f", ..)）比 burn 的手写循环和 Learner：

train_losses             n= 18  max|diff|=1.19e-07
valid_mean               n=  3  max|diff|=2.98e-08
valid_weighted           n=  3  max|diff|=5.96e-08
valid_acc                n=  3  max|diff|=0.00e+00
worst=1.19e-07 tol=1e-06 OK
compare.py exit code: 0
hand vs learner: 18/18 train losses bit-identical
step  torch       burn hand   learner
   0  0.92410517  0.92410517  0.92410517
   1  0.86148870  0.86148876  0.86148876
   5  0.73062712  0.73062712  0.73062712
   6  0.64641875  0.64641887  0.64641887
  11  0.56974900  0.56974906  0.56974906
  12  0.49204874  0.49204874  0.49204874
  17  0.46890214  0.46890217  0.46890217
valid_mean      torch [0.761935, 0.616714, 0.520457]  burn [0.761935, 0.616714, 0.520457]
valid_weighted  torch [0.740081, 0.593879, 0.498393]  burn [0.740081, 0.593879, 0.498393]
valid_acc       torch [75.0, 75.0, 83.333333]  burn [75.0, 75.0, 83.333333]

PyTorch 对 burn：18 步训练 loss 最大差 1.19e-7（第 4 步，0.763 附近两个 f32 ulp），容差 1e-6 通过；验证准确率完全相同。逐步看，差别在 0 到 2 个 ulp 之间来回，最后 6 步里 5 步逐位相同，没有像优化器状态里的 Adam 那样随步数增长。
burn 手写对 Learner：18 步逐位相同。同样的初值、同样顺序的批、同一串调用；Learner 多做的事（发事件、算指标、存文件）都不碰参数。不打乱、单线程加载时才有这个结论，打乱时 Learner 对加载器做的 to_device 会改变顺序，见 burn 卷的优化器。

指标放在哪

PyTorch 的循环里，指标是你自己写的那几行：上面的 s、ws、correct。Learner 里，指标是注册给 SupervisedTraining 的对象，在一个单独的 train-worker 线程上算，每个 epoch 每个指标写一个日志文件（机制见 Learner 与数据）。CliMetricsRenderer 只打印进度（这次 24 行 ProgressSnapshot），数值在最后的摘要表里：

| Split | Metric   | Min.     | Epoch    | Max.     | Epoch    |
|-------|----------|----------|----------|----------|----------|
| Train | Loss     | 0.483    | 3        | 0.818    | 1        |
| Valid | Accuracy | 75.000   | 1        | 83.333   | 3        |
| Valid | Loss     | 0.520    | 3        | 0.762    | 1        |

第 1 个 epoch 验证集的两个日志文件（artifacts/run/valid/epoch-1/，每行是「值,权重」，最后一行是这个 epoch 的汇总）：

== Loss.log
0.696373462677002,1
0.8274973630905151,1
0.7619354128837585,final
== Accuracy.log
81.25,16
62.5,8
75,final

📝NOTE

没有蠢问题

问：摘要表里第 1 个 epoch 的验证 Loss 是 0.762，我在 PyTorch 里按样本数加权算的是 0.740，哪个对？

答：都对，平均的方式不同。LossMetric::update 拿 loss 张量的长度当这批的权重（let [batch_size] = loss.tensor.dims();），交叉熵取 mean 后长度是 1，所以每批权重都是 1，epoch 的值是两批的算术平均：(0.6964 + 0.8275) / 2 = 0.7619，等于上面的 valid_mean。AccuracyMetric::update 用的权重是这批的样本数（16 和 8），(81.25 × 16 + 62.5 × 8) / 24 = 75，按样本加权。最后一批不满时，Learner 的 Loss 和常见的 PyTorch 写法（total += loss.item() * len(xb)）对不上；批大小都相同时两者一致。

checkpoint

PyTorch 里存什么、存成几个文件都由你决定，上面每个 epoch 存一个 .pt，里面是模型和优化器的 state_dict：

2845 epoch-1.pt
2845 epoch-2.pt
2845 epoch-3.pt

with_default_checkpointers() 每个 epoch 存三个 burnpack 文件（模型、优化器、学习率调度器），并按默认策略删旧的：

1292 model-2.bpk
1292 model-3.bpk
512 optim-2.bpk
512 optim-3.bpk
256 scheduler-2.bpk
256 scheduler-3.bpk

第 1 个 epoch 的已经被删掉（「保留最近 2 个」和「验证 loss 最低的那个」都不再要它；规则见 Learner 与数据）。优化器文件只有 512 字节：不带动量的 SGD 没有要存的张量。手写的 burn 循环要存 checkpoint，同样是模型和优化器各存各的（optim.save(..)，见优化器状态和保存与加载）。

训练完的模型

launch 返回 LearningResult，它的 model 字段的文档是「The model with the learned weights, converted to validation mode」，并提醒接着训练前先调 Module::train。实测：

let m = &result.model;
eprintln!("result.model require_grad: {}", m.fc1.weight.val().is_require_grad());
eprintln!("after .train(): {}", m.clone().train().fc1.weight.val().is_require_grad());
eprintln!("ParamId kept: {}", m.fc1.weight.id == model0.fc1.weight.id);

result.model require_grad: false
after .train(): true
ParamId kept: true

ParamId 和传进去的 model0 相同，所以训练完的模型和 checkpoint 里的优化器状态还能对上（原因见优化器状态）。

❗IMPORTANT

动脑筋：不调 .train()，直接拿 result.model 在 autodiff 设备上的一批数据上前向、backward，会发生什么？

let batch = SampleBatcher.batch(samples(0..16), &device);
let model = result.model;
let out = model.forward_classification(batch);
eprintln!("loss is_tracked: {}", out.loss.is_tracked());
let grads = GradientsParams::from_grads(out.loss.backward(), &model);

device 是 Device::flex().autodiff()。先想好再看本页末尾的答案。

开销

同一个二进制连跑 5 次（release，uptime 的 1 分钟负载 8.09，别的实验在并行跑），各取最小值：手写循环 1.24 ms，Learner（含它起线程、写日志、存 checkpoint）56.8 ms。这个任务每步只有几十微秒的计算，Learner 的时间几乎全是这些固定开销；模型和数据大了，这部分就不显眼了。本页不对开销做拆分。

动脑筋答案

loss is_tracked: false
=== PANIC ===
A fatal error happened, you can check the experiment logs here => 'artifacts/run/experiment.log'
=============

thread 'main' (8680819) panicked at .../burn-tensor-0.22.1/src/tensor/api/autodiff.rs:75:9:
Tensor::backward requires a tracked autodiff tensor; call Tensor::autodiff().require_grad() on the source leaf before computing the output
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

前向能跑，但 loss 没有被追踪：result.model 的参数已经去掉 autodiff、不要梯度，输入虽然在 autodiff 设备上也不要梯度，整张图里没有一个叶子要求导。
backward panic，原文和backward 与梯度里对梯度再求导时的那条相同。改法是先 result.model.train()。
中间那三行 === PANIC === 不是 burn-tensor 打印的。launch 装应用日志时用 std::panic::set_hook 挂了一个钩子（learner/application_logger.rs），先把 panic 写进 experiment.log，再提示去看日志。训练已经结束，钩子还在，这个进程后面任何 panic 都会先打印这三行。

小结

burn 的手写循环和 PyTorch 一行对一行：model.train() / model.eval() + no_grad 换成训练时用 model、验证时用 model.valid() 的副本；验证数据放在 device.inner() 上。
同一份数据和初值，PyTorch 和 burn 的手写循环 18 步 loss 最大差 1.19e-7，验证准确率相同；burn 手写循环和 Learner 18 步逐位相同（不打乱、单线程加载）。
Learner 的指标在独立线程上算、写进每个 epoch 的日志文件，摘要表只在结束时打印。LossMetric 每批权重是 1（批的算术平均），AccuracyMetric 按样本数加权；最后一批不满时 Loss 和按样本加权的 PyTorch 写法不同（0.762 对 0.740）。
checkpoint：PyTorch 自己存一个字典；Learner 每个 epoch 存模型、优化器、调度器三个文件，按策略删旧的。
LearningResult.model 是 valid() 过的，ParamId 不变；接着训练前要 .train()，否则 backward panic，Learner 装的 panic 钩子在训练结束后仍然生效。
这个小任务上 Learner 整体 56.8 ms，手写循环 1.24 ms（5 次取最小）。