定义模型

主线位置：训练主线的起点，forward 之前把模型建出来：结构、构造参数、参数放在哪个设备上。

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

flowchart TB
  subgraph pt ["PyTorch"]
    cls["class Mlp nn.Module"] --> init["__init__ 赋属性"]
    init --> reg["setattr 登记参数"]
  end
  subgraph bn ["burn"]
    cfg["MlpConfig new"] -->|"init device"| st["Mlp 结构体"]
    der["derive Module"] -.->|"编译期生成遍历"| st
  end
  reg --> fwd["forward"]
  st --> fwd

实验代码：PyTorch 侧 define.py、register.py，burn 侧 crate module-t 的 bin define、register（burn =0.22.1，feature std、flex、ndarray、autodiff、optim，debug），HashMap 那两个程序在 crate module-err。平台 M2 Max、macOS 27、rustc 1.99.0。

一个结构体，不是一个子类

主线的两层 MLP，两边各写一遍：

class Mlp(nn.Module):
  def __init__(self, d_in, d_hidden, d_out):
    super().__init__()
    self.fc1 = nn.Linear(d_in, d_hidden)
    self.fc2 = nn.Linear(d_hidden, d_out)

  def forward(self, x):
    return self.fc2(torch.relu(self.fc1(x)))

model = Mlp(4, 8, 3)
print(model)

#[derive(Config, Debug)]
pub struct MlpConfig {
  pub d_in: usize,
  pub d_hidden: usize,
  pub d_out: usize,
}

#[derive(Module, Debug)]
pub struct Mlp {
  pub fc1: Linear,
  pub fc2: Linear,
  pub relu: Relu,
}

impl MlpConfig {
  pub fn init(&self, device: &Device) -> Mlp {
    Mlp {
      fc1: LinearConfig::new(self.d_in, self.d_hidden).init(device),
      fc2: LinearConfig::new(self.d_hidden, self.d_out).init(device),
      relu: Relu::new(),
    }
  }
}

impl Mlp {
  pub fn forward(&self, x: Tensor<2>) -> Tensor<2> {
    self.fc2.forward(self.relu.forward(self.fc1.forward(x)))
  }
}

let device = Device::flex();
let config = MlpConfig::new(4, 8, 3);
println!("{config}");
let model = config.init(&device);
println!("{model}");

Mlp(
  (fc1): Linear(in_features=4, out_features=8, bias=True)
  (fc2): Linear(in_features=8, out_features=3, bias=True)
)

{
  "d_in": 4,
  "d_hidden": 8,
  "d_out": 3
}
Mlp {
  fc1: Linear {d_input: 4, d_output: 8, bias: true, params: 40}
  fc2: Linear {d_input: 8, d_output: 3, bias: true, params: 27}
  relu: Relu
  params: 67
}

概念上是同一件事：一个模型由若干子模块组成，forward 把它们串起来。写法上有三处不同。

基类换成派生宏。 PyTorch 在运行期给属性赋值时，把 nn.Module 和 nn.Parameter 类型的值登记下来（下文 register.py 的输出能看出只认这两类）。burn 的 Module 是一个 trait，#[derive(Module)] 在编译期按字段生成 visit、map、to_device、valid、train 这些方法，每个方法逐字段递归。没有登记表，模块树的形状写死在生成的代码里（展开后的代码见 burn 卷 Module 派生）。Clone 也由这个宏生成，再写 derive(Clone) 是 E0119；Debug 要自己写。
构造参数单独成型。 PyTorch 的构造参数是 __init__ 的形参，用完就没了，state_dict() 里只有张量。burn 的习惯是把它们写成一个 #[derive(Config)] 结构体：宏生成 new（必填字段）、with_*（带默认值的字段）、serde 序列化，Display 打印成上面那段 JSON。init(&device) 由你自己写，burn-nn 里的层都按这个约定提供 XxxConfig::init。配置存成 JSON、权重存成记录，加载时先 Config::load 再 init，再把权重读进去（见 burn 卷 Config）。
设备在建的时候给。 nn.Linear(4, 8) 当场在 CPU 上分配并随机初始化，换设备要事后 .to(device)。LinearConfig::init(&device) 拿到设备，但不分配：参数是懒的，第一次读值时才在那个设备上抽随机数（Module 派生的「懒初始化」）。所以 Mlp 不带后端类型参数，同一个类型放在 Flex 上或 Metal 上都行，由 init 时传的 Device 决定。

打印的模块树两边信息差不多：PyTorch 是 repr，burn 是派生宏生成的 Display，每层多一个参数个数，末尾是总数 67。

Linear 的权重是转置的

print("fc1.weight", list(model.fc1.weight.shape), "fc1.bias", list(model.fc1.bias.shape))
print("device", model.fc1.weight.device)

println!("fc1.weight {:?} fc1.bias {:?}", model.fc1.weight.dims(), model.fc1.bias.as_ref().map(|b| b.dims()));
println!("devices {:?}", burn::module::Module::devices(&model));

fc1.weight [8, 4] fc1.bias [8]
device cpu

fc1.weight [4, 8] fc1.bias Some([8])
devices [Device<Flex(Cpu)>]

PyTorch 的 nn.Linear 权重是 [out_features, in_features]，算 x @ W.T + b；burn 的 Linear 是 [d_input, d_output]，Linear::forward 直接 x @ W + b。偏置是 Option<Param<Tensor<1>>>，所以要 as_ref() 才能看。

同一组权重喂给两边，只要 burn 侧转置一次，前向就对得上。A1 的做法：按 PyTorch 的 [out, in] 下标生成数值，from_data 成 [out, in] 再 transpose()，用 Param::from_tensor 换掉懒参数：

model.fc1.weight = Param::from_tensor(weight(8, 4, f64::sin, device)); // [8, 4] -> transpose -> [4, 8]
model.fc1.bias = Some(Param::from_tensor(bias(8, 0.01, device)));

define 两边都对 A1 的 4 个样本做一次前向，存成 JSON 再对比：

logits                   n= 12  max|diff|=1.49e-08
worst=1.49e-08 tol=1e-06 OK

12 个 logits 最大差 1.49e-8，容差 1e-6 通过；不是 0 是因为矩阵乘的累加顺序不同（见怎么读这一卷的「「对得上」是什么意思」）。直接从 safetensors 读 PyTorch 的权重、不做转置时，非方阵报 ShapeMismatch，方阵静默读反，见 burn 卷 ShapeMismatch expected [4, 3] found [3, 4]；从文件导入 PyTorch 权重时由适配器负责转置，见导入 PyTorch 权重。

默认初始化两边是同一个分布：nn.Linear 的文档字符串写的是权重和偏置都「initialized from U(-sqrt(k), sqrt(k))」，k = 1/in_features；burn 的默认值 KaimingUniform { gain: 1/sqrt(3) } 代进去也是 U(-1/sqrt(in), 1/sqrt(in))（见 burn 卷 nn 模块的「初始化器」）。Linear(256, 64) 各抽一次：

torch.manual_seed(1)
big = torch.nn.Linear(256, 64)
print(f"Linear(256, 64) |weight| max {big.weight.abs().max().item():.5f}  |bias| max {big.bias.abs().max().item():.5f}  bound {1 / 256 ** 0.5:.5f}")

device.seed(1);
let big = burn::nn::LinearConfig::new(256, 64).init(&device);
let wmax: f32 = big.weight.val().abs().max().into_scalar();
let bmax: f32 = big.bias.as_ref().map(|b| b.val().abs().max().into_scalar()).unwrap_or(0.0);
println!("Linear(256, 64) |weight| max {wmax:.5}  |bias| max {bmax:.5}  bound {:.5}", 1.0 / 256f32.sqrt());

Linear(256, 64) |weight| max 0.06249  |bias| max 0.06056  bound 0.06250

Linear(256, 64) |weight| max 0.06250  |bias| max 0.06228  bound 0.06250

分布一样，数值不一样：随机数生成器不同，同一个种子抽出来的不是同一组数。要两边起点相同，只能像 A1 那样把权重写死，或者从一边导入另一边。

登记：看属性的值，还是看字段的类型

PyTorch 的登记发生在赋值那一刻，只认 nn.Module、nn.Parameter 和 register_buffer。两个常见的漏法：

class PlainList(nn.Module):
  def __init__(self):
    super().__init__()
    self.layers = [nn.Linear(4, 8), nn.Linear(8, 3)]
    self.scale = torch.ones(3)

class Registered(nn.Module):
  def __init__(self):
    super().__init__()
    self.layers = nn.ModuleList([nn.Linear(4, 8), nn.Linear(8, 3)])
    self.register_buffer("scale", torch.ones(3))

for m in (PlainList(), Registered()):
  n = sum(p.numel() for p in m.parameters())
  print(f"{type(m).__name__:10} params={n:3} state_dict={list(m.state_dict().keys())}")

m = PlainList().to(torch.float64)
print("PlainList after .to(float64):", m.layers[0].weight.dtype, m.scale.dtype)

PlainList  params=  0 state_dict=[]
Registered params= 67 state_dict=['scale', 'layers.0.weight', 'layers.0.bias', 'layers.1.weight', 'layers.1.bias']
PlainList after .to(float64): torch.float32 torch.float32

普通 list 里的两层不在 parameters() 里，优化器看不到，state_dict() 也不存，而且不报错；.to(torch.float64) 也不碰它们，转完还是 float32。普通张量属性同样不登记，要 register_buffer。

burn 的派生宏看的是字段的类型，不看值。Vec<Linear> 的元素类型不是基本类型，宏就把它当子模块，burn-core 给 Vec<T> 写好了 Module 实现，逐个元素递归，路径用下标：

#[derive(Module, Debug)]
struct Stack {
  layers: Vec<Linear>,
  scale: Tensor<1>,
}

let device = Device::flex().autodiff();
let stack = Stack {
  layers: vec![LinearConfig::new(4, 8).init(&device), LinearConfig::new(8, 3).init(&device)],
  scale: Tensor::ones([3], &device),
};
let mut named = Named::default(); // 打印路径的访问者，见下一页
stack.visit(&mut named);

layers.0.weight  [4, 8]   require_grad=true  id=3pda55c6k2qoq
layers.0.bias    [8]      require_grad=true  id=6sk6lf0dni99m
layers.1.weight  [8, 3]   require_grad=true  id=o5t5kda3b6604
layers.1.bias    [3]      require_grad=true  id=vppvga9mpsusi
...

Vec<Linear> 就是 nn.ModuleList，路径 layers.0.weight 也和 PyTorch 一样。Option<T>、[T; N]、2 到 10 元的元组也都是容器模块。裸的 Tensor<1> 字段是另一回事，答案在本页末尾的动脑筋里。

nn.ModuleDict 没有对应。把子模块放进 HashMap<String, Linear>，派生宏同样把它当子模块，但 HashMap 没有 Module 实现（crate module-err，bin dict_err，debug；rustc 报 10 个错误、展开 3 个，节选第一个）：

error[E0277]: the trait bound `HashMap<String, burn::nn::Linear>: Module` is not satisfied
 --> src/bin/dict_err.rs:5:10
  |
5 | #[derive(Module, Debug)]
  |          ^^^^^^ the trait `Module` is not implemented for `HashMap<String, burn::nn::Linear>`
  |
  = note: If `HashMap<String, burn::nn::Linear>` is used as configuration or runtime state in a `#[derive(Module)]` type, mark that field with `#[module(skip)]`.
...

这句 note 是写给配置字段的（见 burn 卷 the trait bound X: Module is not satisfied）。照着加上 #[module(skip)] 能编译，但里面的层就成了常量（bin dict_skip）：

#[derive(Module, Debug)]
struct Heads {
  #[module(skip)]
  heads: HashMap<String, Linear>,
}

num_params = 0 (the Linear alone has 27)

和 PyTorch 的普通列表一样，参数从模型里消失了，优化器不会更新它们。按名字取子模块，改用 Vec<Linear> 加一张名字到下标的表，或者每个头一个具名字段。

📝NOTE

没有蠢问题

问：没有 nn.Sequential 吗？我只想把几层串起来。

答： burn-nn 0.22.1 里没有 Sequential（源码里搜不到这个名字），要像上面的 Mlp 那样写一个结构体和 forward。原因在类型上：forward 的签名是 Tensor<D> 进、Tensor<D> 出，秩在类型里，Linear、Conv2d、Relu 的 forward 又是各自类型上的普通方法，不是同一个 trait，Vec<T> 只能装同一种层。层数可变、类型相同（一叠 Linear），用 Vec<Linear> 加一个 for 循环；类型不同又想放进一个容器，派生宏支持枚举，每个变体一个字段（enum Block { Small(Linear), Large(Head) }），forward 里 match。

forward 是普通方法

PyTorch 里写 model(x) 而不是 model.forward(x)，Module.forward 的文档字符串说前者「takes care of running the registered hooks while the latter silently ignores them」。burn 的 forward 是 impl Mlp 里一个普通的方法，名字和签名都是约定，Module trait 里没有它，调用就是 model.forward(x)，中间没有 hook。签名里写明了秩（Tensor<2>），传错秩是编译错误（见创建与形状）。forward 拿 &self，参数用 self.fc1.weight.val() 取出来，val() 克隆的是句柄，不复制数据（见运算消费张量）。

❗IMPORTANT

动脑筋：裸张量字段 scale 进不进记录？

上面的 Stack 里，scale: Tensor<1> 没有包成 Param。PyTorch 的 register_buffer("scale", ...) 让它进了 state_dict（5 个键）。register 接着打印：

println!("num_params={} record entries={}", stack.num_params(), stack.clone().into_record().len());
let valid = stack.valid();
println!("after valid(): scale is_autodiff={}", valid.scale.is_autodiff());

burn 侧 into_record() 有几条？valid() 之后，scale 还在 autodiff 设备上吗？

动脑筋答案

num_params=67 record entries=4
after valid(): scale is_autodiff=false

记录只有 4 条，scale 不在里面，num_params 也不算它。裸 Tensor 字段的 Module 实现里 visit 和 map 什么也不做，源码上方的注释是 TODO: tensor record should persist。访问者看不到它，优化器和记录也就看不到它。
valid() 之后它不在 autodiff 设备上了：to_device、valid、train 这些方法对它是实现了的，它跟着模型搬设备、换状态。
所以裸张量字段像 PyTorch 的 register_buffer(..., persistent=False)（实测这种 buffer 在 named_buffers() 里、不在 state_dict() 里）：跟着模型走，不存盘。要存盘又不训练，包成 Param 再 set_require_grad(false)（register 最后一段：只有这一个字段的模块记录有 1 条，valid().train() 之后仍不要梯度；它会算进 num_params）；BatchNorm 的运行统计用的是另一种类型 RunningState，见参数与遍历。

小结

nn.Module 子类对 #[derive(Module)] 结构体：派生宏在编译期按字段生成遍历，没有运行期登记表；构造参数放进 #[derive(Config)]，init(&device) 建模型，参数是懒的。
Linear 权重 PyTorch 是 [out, in]，burn 是 [in, out]；转置一次，A1 的前向 12 个数最大差 1.49e-8。默认初始化分布相同，随机数不同。
PyTorch 的普通列表、普通张量属性不登记；burn 的 Vec<Linear> 就是 ModuleList，裸 Tensor 字段跟着模型走但不进记录。
HashMap 装不了子模块，E0277 的 note 建议的 #[module(skip)] 会让参数消失；没有 Sequential，forward 手写。