模块与训练

这一部分把训练主线拆开讲。前面几部分讲的都是「一个张量算子」和「一次 backward」；到这里，对象变成一整个模型：它由哪些参数组成、怎么建、怎么前向、梯度怎么回到参数上、参数怎么存盘，以及谁来跑这个循环。涉及的 crate 有 burn-core（Module、Param、Config、记录）、burn-derive（两个派生宏）、burn-nn（现成的层）、burn-optim（优化器）、burn-store（文件格式）和 burn-train（Learner）。

0.22 里这一层最大的变化是没有后端类型参数了：模型写成 struct Mlp { l1: Linear, .. }，不再是 Mlp<B: Backend>；训练态和验证态是同一个类型的两种状态，而不是 Autodiff<B> 和 B 上的两个类型。一个手写的完整训练步见一个训练步，本部分各页讲它每一步背后的机制。

1. Module 派生：Module trait 的八个必需方法由 #[derive(Module)] 逐字段递归生成；字段按语法分成子模块、常量和 #[module(skip)]；Param 带一个随机的 ParamId、懒初始化的值和单独记着的训练状态；ModuleVisitor / ModuleMapper 拼出参数路径；valid / train 在同一个类型上切换；to_device 和 fork 对已初始化的参数结果不同。
2. Config：#[derive(Config)] 生成 new、with_*、影子结构体做 serde 和打印 JSON 的 Display；默认值只在 new 里生效；记录里只有张量，所以模块要用 XxxConfig::init(device) 两段式建，重建模型需要配置和记录两样。
3. nn 模块：Linear、Conv2d、LayerNorm、Dropout 都是参数加一次函数式调用，函数式算子经 Dispatch 落到 ModuleOps；linear 在 0.22.1 里绕过了 Autodiff 的专用反向；初始化器是懒的；Dropout 只在 autodiff 设备上生效。
4. 优化器：逐张量的 Optimizer trait 和类型擦除的 ModuleOptimizer（AdamConfig::init() 返回的就是它）；GradientsParams 把按节点 id 存的梯度换成按 ParamId 存；Adam 的状态按参数 id 保存。
5. 存储与加载：into_record / save_file 存成 burnpack；burn-store 读 safetensors 和 PyTorch 文件，键名重映射和权重转置。
6. Learner 与数据：burn-train 的主循环、指标和 checkpoint，DataLoader 怎么分批和多线程加载。

flowchart TB
  cfg["XxxConfig"] -->|init device| mod["Module 和 Param"]
  mod -->|forward| ops["nn 层到 ModuleOps"]
  ops --> bw["backward 得 Gradients"]
  bw --> gp["GradientsParams 按 ParamId"]
  gp --> opt["ModuleOptimizer step"]
  opt -->|map 换参数| mod
  mod --> rec["记录 save_file"]
  learner["Learner"] -.->|循环| ops

前三页在「一个模型是什么」这一侧：它的结构、它的配置、它的前向；后三页在「怎么训练它」这一侧：更新参数、存取参数、跑循环。两侧的接缝是 Module 派生出来的遍历：前向和 backward 只认计算图里的节点 id，优化器认 ParamId，记录按遍历拼出的参数路径（layers.0.weight）存取，三种标识靠同一次遍历对上。