Learner 与数据

主线位置：训练主线的最外层。一个 step 里的前向、反向、GradientsParams、ModuleOptimizer::step 已经在一个训练步和优化器讲过；本页讲 burn-train 怎样把这个 step 放进 epoch 循环，配上数据加载、验证、指标、渲染和 checkpoint，以及 burn-core 的 DataLoader 怎样把数据集切成批。

sequenceDiagram
  participant U as 用户 main
  participant S as SupervisedTraining
  participant E as 单设备 epoch
  participant L as Learner
  participant W as train-worker 线程
  U->>S: launch learner
  S->>E: fit 每个 epoch
  loop 每一批
    E->>L: train_step 前向和反向
    E->>L: lr_step 然后 optimizer_step
    E->>W: ProcessedItem 事件
    W->>W: sync 算指标 写日志
  end
  E->>E: 验证 epoch 用 valid 模型
  E->>S: 按策略存 checkpoint
  S-->>U: LearningResult

实验 crate 是 train-learner（从模板 base 复制，Flex，debug），和优化器页的实验是同一个程序：两层 MLP 回归，200 个训练样本、64 个验证样本，每批 32，3 个 epoch。平台 M2 Max、macOS 27、rustc 1.99.0，输出存在 train-learner/out/。

用户写的部分：TrainStep 和 InferenceStep

模型要被 Learner 训练，得实现两个 trait。TrainStep：

pub trait TrainStep {
    /// Type of input for a step of the training stage.
    type Input: Send + 'static;
    /// Type of output for a step of the training stage.
    type Output: ItemLazy + 'static;
    // ...
    fn step(&self, item: Self::Input) -> TrainOutput<Self::Output>;
    // ...
    fn optimize(
        self,
        optim: &mut ModuleOptimizer,
        lr_module: ModuleLearningRate,
        grads: GradientsParams,
    ) -> Self
    where
        Self: Module + Sized,
    {
        optim.step(lr_module, self, grads)
    }
    // ...
}

step 做前向和反向，交回一个 TrainOutput：TrainOutput::new(self, loss.backward(), item) 当场调 GradientsParams::from_grads，把梯度换成按 ParamId 存。optimize 有默认实现，想在更新前后对模型做点别的（比如裁剪权重）可以覆盖。验证用的 InferenceStep 只有一个 step(&self, item) -> Output，没有反向。0.20 之前验证用的 trait 叫 ValidStep<VI, VO>，输入输出是类型参数；#4283（「Clean learning api」，v0.20.0）改成了关联类型的 InferenceStep。

LearnerModel 是把要求捆在一起的空 trait：TrainStep + InferenceStep + Module + Display + 'static，#[derive(Module)] 已经实现了 Display（结尾的训练摘要就用它打印模型结构）。

Output 要实现 ItemLazy，并对想用的指标实现 Adaptor<指标的输入>。burn 自带回归、分类、多标签分类和序列几种输出；本页用 RegressionOutput（loss、output、targets 三个张量），它的 sync 只 flush 一下设备，再对三个张量调 without_autodiff()：交给指标线程的张量不再挂着计算图，也不在这里读回数据。

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

  fn step(&self, batch: Batch) -> TrainOutput<RegressionOutput> {
    let out = self.forward_regression(batch);
    LEARNER_LOSSES.lock().unwrap().push(out.loss.clone().into_scalar());
    TrainOutput::new(self, out.loss.backward(), out)
  }
}

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

  fn step(&self, batch: Batch) -> RegressionOutput {
    self.forward_regression(batch)
  }
}

（LEARNER_LOSSES 是给优化器页对比用的，不是必需的。）

组装：Learner 和 SupervisedTraining

Learner 只有四个字段：模型、ModuleOptimizer、ModuleLrScheduler、当前学习率。Learner::new(model, optim, lr_scheduler) 的第三个参数收任何 LrScheduler，传 1e-2 就是常数学习率。训练的其余部分由 SupervisedTraining 配置，new(directory, 训练加载器, 验证加载器) 时的默认值：

• num_epochs 是 1。
• 应用日志写到 directory/experiment.log（with_application_logger(None) 关掉）。
• checkpoint 策略是「保留最近 2 个」和「验证集 loss 均值最低的那个」的组合；但默认没有 checkpointer，要 with_default_checkpointers() 才真的存。
• 渲染器、指标日志器都没设，launch 时补默认值。

本页的配置：

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

launch 依次：装应用日志；没指定渲染器就取 default_renderer；没有指标日志器就注册一个写到 directory 的 FileMetricLogger；把指标、渲染器和事件存储包成 FullEventProcessorTraining，再包一层 AsyncProcessorTraining；建 checkpointer；最后选训练策略。没指定策略时用单设备，设备取模型第一个参数所在的设备，经过 autodiff_device：不是 .autodiff() 的设备会被套上 autodiff，.gradient_checkpointing() 打开时再加检查点。所以模型建在普通的 Device::flex() 上也能训练，fit 开头的 learner.fork(&device) 会把它搬到 autodiff 设备。多设备、DDP 和自定义策略通过 with_training_strategy 选，本页不展开。

一个 epoch

SingleDeviceTrainingStrategy::fit 先把训练加载器 to_device 到 autodiff 设备，验证加载器 to_device 到 device.inner()（不带 autodiff 的同一设备），然后每个 epoch：训练一遍、验证一遍、发 EndEpoch；有 checkpointer 或早停时先 flush 事件处理器，让指标算完，再按策略存 checkpoint、问早停。训练的那一遍是 SingleDeviceTrainEpoch::run：

        while let Some(item) = iterator.next() {
            // ...
            iteration += 1;
            log::info!("Iteration {iteration}");

            let progress = iterator.progress();
            let item = learner.train_step(item);

            match self.grad_accumulation {
                Some(accumulation) => {
                    accumulator.accumulate(&learner.model(), item.grads);
                    accumulation_current += 1;

                    if accumulation <= accumulation_current || progress.is_completed() {
                        let grads = accumulator.grads();

                        learner.lr_step();
                        learner.optimizer_step(grads);
                        accumulation_current = 0;
                    }
                }
                None => {
                    learner.lr_step();
                    learner.optimizer_step(item.grads);
                }
            }

            let item = TrainingItem::new(
                item.item,
                progress,
                Some(iteration),
                Some(learner.lr_current()),
            );

            interrupter.fail_on_error(processor.process_train(LearnerEvent::ProcessedItem(item)));
            // ...
        }

每批：train_step（用户的 TrainStep::step）、lr_step（调度器走一步，得到这一步的学习率）、optimizer_step（TrainStep::optimize）、把输出连同进度和学习率作为事件交出去。梯度累积时（grads_accumulation(n)），梯度先在 GradientsAccumulator 里按 ParamId 相加，攒够 n 批或到了 epoch 末尾才调一次调度器和优化器。验证那一遍（SingleDeviceValidEpoch::run）先取 learner.model().valid()，逐批调 InferenceStep::step，同样交出事件。valid() 去掉参数的 autodiff（见 Module 派生），验证时不建计算图。

训练结束后 train 返回的 LearningResult.model 也是 valid() 过的；它的文档提醒，要接着训练得先 .train()。设备上的错误不 panic，记在 Interrupter 里，从 LearningResult.error 取。

事件、指标和渲染器

训练循环不直接算指标，它只发 LearnerEvent：Start、StartSplit、ProcessedItem、EndSplit、EndEpoch、End。AsyncProcessorTraining 起一个名为 train-worker 的线程（start），事件经一个容量为 1 的 async_channel 送过去，所以指标计算和下一批的训练并行；通道里只能排一个事件，指标线程跟不上时训练循环会在发送处等它。那个线程上的 FullEventProcessorTraining 处理 ProcessedItem：先 item.sync()，再更新每个指标，得到的条目交给渲染器、写进事件存储；事件存储的日志器把每个值写成文件。EndSplit 时算 epoch 的汇总值。

LossMetric 的 update 读 loss.mean() 回来，以 loss 张量的长度作为这批的权重：let [batch_size] = loss.tensor.dims();。MseLoss 配 Reduction::Mean 得到的是长度 1 的张量，所以每批权重都是 1，epoch 的值是各批 loss 的算术平均，不按样本数加权。

渲染器由 default_renderer 决定：开了 tui feature（burn-train 的默认 feature，burn 开默认 feature 时会带上）并且 stdout 是终端，用全屏的 TUI；否则用 CliMetricsRenderer。不想要 TUI 有三种办法：像本页这样 .renderer(Box::new(CliMetricsRenderer::new()))；把输出重定向到文件或管道；burn 关掉默认 feature。CliMetricsRenderer 只打印进度，它的 update_train / update_valid 是空函数，指标的值不打印；要看数，看 .summary() 在结束时打印的表，或者读日志文件，或者自己实现 MetricsRenderer 和它要求的四个 trait（训练、评估两侧的渲染和进度日志；burn 仓库的 examples/custom-renderer 是一个模板）。

checkpoint

with_default_checkpointers() 建三个 FileCheckpointer，目录都是 directory/checkpoint，名字分别是 model、optim、scheduler，文件名是 {名字}-{epoch}.bpk。三者存的分别是 model.into_record()、optim.to_record()、调度器的记录，都是 burnpack 文件（格式见存储与加载）；各包在一个 AsyncCheckpointer 里，在后台线程写盘。LearningCheckpointer::checkpoint 每个 epoch 末问一次策略，策略交回一串 Save / Delete(epoch)。

默认策略是一个 ComposedCheckpointingStrategy，里面两个：KeepLastNCheckpoints::new(2)（每次都存，并要求删掉 epoch - 2）和 MetricCheckpointingStrategy（验证集 Loss 的均值最低的 epoch 要留着，最佳 epoch 变了就要求删掉旧的）。组合的规则是：任一子策略要存就存；一个 epoch 只有在所有子策略都不要它时才删。续训用 .checkpoint(n)：launch 先 load_checkpoint(learner, n) 读回三个文件，再从第 n + 1 个 epoch 开始。

DataLoader

数据这一侧在 burn-core/src/data/dataloader/。三个角色：

• Dataset<I>（burn-dataset）：能按下标取样本。本页用 InMemDataset::new(Vec<Sample>)。
• Batcher<I, O>：fn batch(&self, items: Vec<I>, device: &Device) -> O，把一批样本拼成张量，建在加载器给的设备上。
• DataLoader<O>：iter() 出批，num_items()，to_device(&Device) 返回一个改了设备的新加载器，slice(start, end) 切一段（多设备训练用）。

impl Batcher<Sample, Batch> for SampleBatcher {
  fn batch(&self, items: Vec<Sample>, device: &Device) -> Batch {
    let n = items.len();
    let xs: Vec<f32> = items.iter().flat_map(|s| s.x).collect();
    let ys: Vec<f32> = items.iter().map(|s| s.y).collect();
    Batch {
      x: Tensor::<2>::from_data(TensorData::new(xs, [n, 4]), device),
      y: Tensor::<2>::from_data(TensorData::new(ys, [n, 1]), device),
    }
  }
}

DataLoaderBuilder 有四个选项：batch_size(n)（FixBatchStrategy，不设时每批 1 个）、shuffle(seed)、num_workers(n)、set_device(device)（不设时是 Device::default()；交给 Learner 的加载器会被 to_device 改掉，自己迭代时要设）。build 在 num_workers > 0 时建 MultiThreadDataLoader，否则建 BatchDataLoader。

BatchDataLoader 在调用方线程上按顺序取样本、攒够一批就调 Batcher，不足一批的尾巴在最后单独成批（200 个样本、每批 32：六批 32 加一批 8）。打乱时，每次 iter() 用自己的 StdRng 生成一个 ShuffledDataset，随机数发生器继续往前走，所以每个 epoch 顺序不同、整个过程可复现。to_device 和 slice 不复制随机数发生器的状态，而是 rng.fork() 派生一个新的：同一个种子，经过一次 to_device 的加载器和原来的打乱顺序不同。优化器页里手写循环和 Learner 加了 .shuffle(7) 就对不上，原因正是单设备策略对训练加载器做的那次 to_device。

MultiThreadDataLoader 第一次 iter() 时把数据集（打乱的话先整体打乱一次）按批切成 num_workers 段，每段一个 BatchDataLoader，各有一个名为 dataloader-{index} 的线程；批通过一个容量 100 的同步通道送回来，谁先做完谁先送。所以多线程加载时批的顺序取决于线程调度，不打乱也不保证可复现。实测（train-learner - 2，不打乱、2 个线程，连跑两次）：

### args: - 2
step  hand        learner     equal
   0  4.47128010  4.47128010  true
   1  3.73394918  3.73394918  true
   2  6.43713617  6.43713617  true
   6  2.71958590  2.71958590  true
  13  2.14323759  2.12240887  false
  20  1.60553181  1.95099330  false
all 21 losses bit-identical: false
...
### args: - 2
step  hand        learner     equal
   0  4.47128010  5.51330090  false
   1  5.39203453  4.35300779  false

两次运行里，手写循环和 Learner 各自拿到的批次顺序都不一样；单线程时两边逐位相同（见优化器）。

实验：3 个 epoch

train-learner（不带参数，debug）里 Learner 那部分的输出，ProgressSnapshot 行每批一行，这里只留第一个 epoch 的：

Starting training for 3 epochs.
Starting split 'train' with 200 items.
ProgressSnapshot { global: Progress { items_processed: 1, items_total: 3, unit: Some("epochs") }, split: Progress { items_processed: 32, items_total: 200, unit: Some("items") } }
...
ProgressSnapshot { global: Progress { items_processed: 1, items_total: 3, unit: Some("epochs") }, split: Progress { items_processed: 200, items_total: 200, unit: Some("items") } }
Split ended.
Starting split 'valid' with 64 items.
ProgressSnapshot { global: Progress { items_processed: 1, items_total: 3, unit: Some("epochs") }, split: Progress { items_processed: 32, items_total: 64, unit: Some("items") } }
ProgressSnapshot { global: Progress { items_processed: 1, items_total: 3, unit: Some("epochs") }, split: Progress { items_processed: 64, items_total: 64, unit: Some("items") } }
Split ended.
...
training ended.
======================== Learner Summary ========================
Model:
Mlp {
  l1: Linear {d_input: 4, d_output: 16, bias: true, params: 80}
  act: Relu
  l2: Linear {d_input: 16, d_output: 1, bias: true, params: 17}
  params: 97
}
Total Epochs: 3


| Split | Metric | Min.     | Epoch    | Max.     | Epoch    |
|-------|--------|----------|----------|----------|----------|
| Train | Loss   | 2.999    | 3        | 4.484    | 1        |
| Valid | Loss   | 2.829    | 3        | 4.427    | 1        |

learner: done in 109.276541ms

CliMetricsRenderer 打印的全是进度，loss 的数值只出现在最后的摘要表里。directory（这里是 artifacts/run，在实验 crate 目录下）里的文件：

1284 artifacts/run/checkpoint/model-2.bpk
1284 artifacts/run/checkpoint/model-3.bpk
3076 artifacts/run/checkpoint/optim-2.bpk
3076 artifacts/run/checkpoint/optim-3.bpk
256 artifacts/run/checkpoint/scheduler-2.bpk
256 artifacts/run/checkpoint/scheduler-3.bpk
3012 artifacts/run/experiment.log
165 artifacts/run/train/epoch-1/Loss.log
166 artifacts/run/train/epoch-2/Loss.log
167 artifacts/run/train/epoch-3/Loss.log
64 artifacts/run/valid/epoch-1/Loss.log
66 artifacts/run/valid/epoch-2/Loss.log
65 artifacts/run/valid/epoch-3/Loss.log

（artifacts/hand-optim.bpk 是手写循环存的优化器状态，不属于 Learner。）每个指标每个 epoch 一个文件，一行一个值：

== artifacts/run/train/epoch-1/Loss.log
4.471280097961426,1
3.7339491844177246,1
6.437136173248291,1
2.9738106727600098,1
5.077566146850586,1
5.977496147155762,1
2.71958589553833,1
4.484403473990304,final

前 7 行是每批的 loss 和权重 1，最后一行 final 是 7 个数的算术平均 4.4844，就是摘要表里 Train 的 Max。最后一批只有 8 个样本，按样本加权的平均会不一样。experiment.log 是 log crate 的输出（Executing training step for epoch 1、Iteration 1 ……）。常数学习率的 scheduler-*.bpk 没有状态，256 字节里只有头、元数据和对齐。

动脑筋答案

1284 artifacts/run/checkpoint/model-2.bpk
1284 artifacts/run/checkpoint/model-3.bpk

剩第 2、3 个 epoch 的。逐个 epoch 推：

• epoch 1：两个子策略都要存，存 *-1.bpk。
• epoch 2：KeepLastN(2) 要删 2 − 2 = 0，不算；验证 loss 的最佳 epoch 从 1 变成 2，指标策略要删 1。但 KeepLastN 还要 1，不删。存 *-2.bpk。
• epoch 3：KeepLastN 要删 1，指标策略（最佳变成 3）要删 2。epoch 1 两个子策略都不要了，删；epoch 2 KeepLastN 还要，留。存 *-3.bpk。

如果验证 loss 在第 1 个 epoch 最低，第 1 个会一直被指标策略留着，第 3 个 epoch 后会剩 1、2、3 三个。

小结

• 模型实现 TrainStep（前向加反向，TrainOutput::new 里换成 GradientsParams）和 InferenceStep（只前向）；输出类型实现 ItemLazy 和指标的 Adaptor，sync 时去掉 autodiff。
• SupervisedTraining::launch 补默认的渲染器和文件日志器，模型所在设备不是 autodiff 就套上；单设备 epoch 每批 train_step → lr_step → optimizer_step → 发事件，验证用 valid() 的模型。
• 指标在 train-worker 线程上算，写进 directory/{train,valid}/epoch-N/<指标>.log；LossMetric 的 epoch 值是各批的算术平均。CliMetricsRenderer 只打印进度，TUI 只在开了 tui 且 stdout 是终端时出现。
• with_default_checkpointers 在 directory/checkpoint/ 存 model、optim、scheduler 三份 burnpack；默认组合策略只删所有子策略都不要的 epoch。
• DataLoaderBuilder：BatchDataLoader 单线程、顺序可复现，to_device / slice 派生新的打乱种子；MultiThreadDataLoader 按批切段、线程常驻，批的顺序取决于线程调度。