参数与遍历

主线位置：训练主线 backward 和 optimizer step 之间：优化器要找到每个参数和它的梯度，冻结的参数要被跳过。

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

flowchart TB
  m["模型"] -->|"visit"| v["ModuleVisitor"]
  v --> f["visit_float 每个 Param"]
  f --> rg{"is_require_grad"}
  rg -->|"是"| gp["收下梯度"]
  rg -->|"否"| skip["跳过 不更新"]
  gp -->|"按 ParamId"| opt["optimizer step"]

实验代码：PyTorch 侧 params.py、freeze.py，burn 侧 crate module-t 的 bin params、freeze、remap、oldleaf（burn =0.22.1，feature std、flex、ndarray、autodiff、optim，debug）。模型是主线的 MLP，权重同 A1，建在 Device::flex().autodiff() 上。平台 M2 Max、macOS 27、rustc 1.99.0。

named_parameters 对一个访问者

PyTorch 的 named_parameters() 是一个生成器，走的是 __init__ 时登记的 _parameters 和 _modules：

model = a1_model()
for name, p in model.named_parameters():
  print(f"{name:16} {str(list(p.shape)):8} requires_grad={p.requires_grad}")
print("sum numel =", sum(p.numel() for p in model.parameters()))

burn 没有现成的「列出参数名」的方法，Module 上的遍历入口是 visit，它接一个 ModuleVisitor。派生代码进出每个字段时调 enter_module / exit_module，遇到浮点参数调 visit_float，自己维护一个路径栈就能拼出名字：

#[derive(Default)]
pub struct Named {
  path: Vec<String>,
  pub count: usize,
}

impl ModuleVisitor for Named {
  fn enter_module(&mut self, name: &str, _container_type: &str) {
    self.path.push(name.to_string());
  }

  fn exit_module(&mut self, _name: &str, _container_type: &str) {
    self.path.pop();
  }

  fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<D>>) {
    let shape = param.lazy_shape();
    self.count += shape.num_elements();
    println!(
      "{:16} {:8} require_grad={:5} id={}",
      self.path.join("."),
      format!("{:?}", shape.as_slice()),
      param.val().is_require_grad(),
      param.id
    );
  }
}

let model = a1_model(&device);
let mut named = Named::default();
model.visit(&mut named);
println!("visited = {} | num_params = {}", named.count, model.num_params());

fc1.weight       [8, 4]   requires_grad=True
fc1.bias         [8]      requires_grad=True
fc2.weight       [3, 8]   requires_grad=True
fc2.bias         [3]      requires_grad=True
sum numel = 67

fc1.weight       [4, 8]   require_grad=true  id=6d51aj5v5j8so
fc1.bias         [8]      require_grad=true  id=gjlsmvqg8gitq
fc2.weight       [8, 3]   require_grad=true  id=9onfnu8bfufpc
fc2.bias         [3]      require_grad=true  id=met1ro0600qce
visited = 67 | num_params = 67

名字一样。 路径是字段名用 . 连起来，Vec 元素用下标，Option 不占一级（所以是 fc1.bias 而不是 fc1.bias.Some）。这条路径也是存储记录里的键名，从 PyTorch 导入权重时键名能直接对上，靠的就是它（见导入 PyTorch 权重）。
形状转置。 Linear 的权重是 [in, out]（见定义模型）。
总数一样。 num_params 本身也是一个访问者，加的是 lazy_shape().num_elements()，不会触发懒参数的初始化。
参数靠 id 认，不靠对象身份。 PyTorch 的优化器拿着 Parameter 对象本身，梯度挂在对象的 .grad 上。burn 的 Param 带一个 ParamId（随机的 64 位整数，打印成 base32，每次运行都不同），GradientsParams 按它存梯度，优化器按它存动量（机制见 burn 卷 Module 派生的「ParamId」）。

想按 PyTorch 的方式一次拿到 (名字, 张量) 列表，把 println! 换成往 Vec 里推就行；访问者只读（&Param），要改参数用对应的 ModuleMapper。

不是参数的状态

PyTorch 把「随模型保存、但不训练」的张量叫 buffer，和参数分开列：

bn = nn.BatchNorm1d(8)
print("BatchNorm1d(8) params =", sum(p.numel() for p in bn.parameters()),
      "| buffers =", [(n, list(b.shape)) for n, b in bn.named_buffers()])

BatchNorm1d(8) params = 16 | buffers = [('running_mean', [8]), ('running_var', [8]), ('num_batches_tracked', [])]

burn 的 BatchNorm 把运行均值和方差存在 RunningState 里，gamma 和 beta 是普通的 Param。访问者分不分得出这两类？

❗IMPORTANT

动脑筋：BatchNormConfig::new(8).init(&device).num_params() 是多少？

PyTorch 的 parameters() 是 16。burn 侧把同一个 Named 访问者用在 BatchNorm(8) 上，再打印 num_params()（bin params 的后半段）。是 16，还是别的数？

num_batches_tracked 在 burn 里没有对应：BatchNormConfig 只有 num_features、epsilon、momentum 三个字段，没有 PyTorch momentum=None 那种累计平均。从 PyTorch 导入 BatchNorm 时，这个键在 burn 模型里没有对应的位置，多余的键怎么处理见导入 PyTorch 权重。

冻结参数

PyTorch 冻结一部分参数，是对它们调 requires_grad_(False)，优化器照常拿全部参数，冻结的那些没有梯度，step 不动它们。主线 MLP 冻结 fc1，训练一步：

model = a1_model()
model.fc1.requires_grad_(False)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
loss = nn.functional.cross_entropy(model(x), y)
opt.zero_grad()
loss.backward()
print({n: (None if p.grad is None else list(p.grad.shape)) for n, p in model.named_parameters()})
fc1_before = model.fc1.weight.clone()
opt.step()
print("fc1 unchanged:", torch.equal(model.fc1.weight, fc1_before))
model.eval(); model.train()
print("after eval()/train():", {n: p.requires_grad for n, p in model.named_parameters()})

burn 有两种写法：对子模块字段调 Module::no_grad，或者在整个模型上按路径选参数，set_require_grad_group 配 ParamGroup::from_predicate（路径里包含这个子串就算）：

let mut model = a1_model(&device);
match how.as_str() {
  "field" => model.fc1 = model.fc1.no_grad(),
  "group" => model = model.set_require_grad_group(ParamGroup::from_predicate("fc1"), false),
  _ => unreachable!("field or group"),
}
println!("{how}: require_grad [fc1, fc2] = {:?}", require(&model));

let grads = ce.forward(model.forward(x), y).backward();
println!("fc1 grads {:?} | fc2 grads {:?}", grad_shape(&model.fc1, &grads), grad_shape(&model.fc2, &grads));
let grads = GradientsParams::from_grads(grads, &model);
println!("GradientsParams::len() = {}", grads.len());
let fc1_before = model.fc1.weight.val();
let mut optim = SgdConfig::new().init();
model = optim.step(0.1, model, grads);
let same: bool = fc1_before.equal(model.fc1.weight.val()).all().into_scalar();
println!("fc1 unchanged: {same}");
println!("after valid().train(): require_grad [fc1, fc2] = {:?}", require(&model.valid().train()));

{'fc1.weight': None, 'fc1.bias': None, 'fc2.weight': [3, 8], 'fc2.bias': [3]}
fc1 unchanged: True
after eval()/train(): {'fc1.weight': False, 'fc1.bias': False, 'fc2.weight': True, 'fc2.bias': True}

field: require_grad [fc1, fc2] = [false, true]
fc1 grads (None, None) | fc2 grads (Some([8, 3]), Some([3]))
GradientsParams::len() = 2
fc1 unchanged: true
after valid().train(): require_grad [fc1, fc2] = [false, true]

group 那次的输出和 field 逐行相同。更新后的四组参数两边对比（两种写法各对比一次，结果相同）：

params_after.fc1.weight  n= 32  max|diff|=0.00e+00
params_after.fc1.bias    n=  8  max|diff|=3.73e-09
params_after.fc2.weight  n= 24  max|diff|=1.49e-08
params_after.fc2.bias    n=  3  max|diff|=9.31e-10
worst=1.49e-08 tol=1e-06 OK

冻结的参数两边都没有梯度：PyTorch 是 .grad is None，burn 是 grad(&grads) 返回 None，GradientsParams 只收了 fc2 的 2 个。GradientsParams::from_grads 的访问者在 visit_float 里先看 param.is_require_grad()，不要梯度的直接跳过，优化器就不碰它。
fc1.bias 冻结了还有 3.73e-9 的差：这是初始值本身的差。PyTorch 侧的 0.01 * o 先在 f64 里算再转 f32，burn 侧在 f32 里乘，最后一位不同（差在第 6 个元素 0.05 上，0.05 附近 f32 的一个 ulp 正好是 3.73e-9）；它没参与更新，差就原样留着。
冻结两边都是持久的：PyTorch 的 eval() / train() 不碰 requires_grad；burn 的 no_grad 改的是参数里记着的「配置上要不要训练」（is_active），valid().train() 恢复的正是它，所以 fc1 仍然不要梯度。解冻用 set_require_grad(true) 或 unfreeze()。
写法上的差别：PyTorch 原地改，model.fc1 还是那个对象；burn 的 no_grad 按值拿走子模块、返回新的，要赋回字段（model.fc1 = model.fc1.no_grad()），整模型的版本要赋回 model。

no_grad 只管梯度。同名的 PyTorch torch.no_grad() 是关掉记图的作用域，burn 里对应的是 valid() 和不带 autodiff 的设备（见不求导的区域）。还有一个 freeze：参数不要梯度，同时关掉 dropout、BatchNorm 这些层的训练行为，相当于 PyTorch 里对子模块 requires_grad_(False) 再 eval()，见训练和推理模式。

📝NOTE

冷知识：no_grad 曾经也关 dropout

PR #5498（「feat(module): generalize Param with module-owned flags」）给 Dropout、BatchNorm 加上 training: Param<Flag> 时，Module::no_grad 的文档是「Each tensor in the module tree will not require grad, and each layer with no tensor of its own — dropout, noise, randomized activations — stops behaving as it does during training.」Dropout 字段的文档写的是「Cleared by no_grad and matching freeze_group traversals」。一周后的 PR #5537（「feat(module): separate gradient control from module freezing」）把两件事拆开：no_grad 改成「This is equivalent to set_require_grad(false). Module-owned control flags are left unchanged」，并写明「A later train keeps gradients disabled; it does not undo no_grad」；关掉训练行为的部分交给新的 freeze，Dropout 文档里的 no_grad 换成了 freeze。同一个 PR 的提交列表里还有一条「Fix RunningState param: not an optimizer target」。两个 PR 首次发布都在 v0.22.0-pre.4，所以「no_grad 关 dropout」只存在于两次提交之间，没有进过任何发布版本。

改参数的值

PyTorch 改参数的值（剪枝、缩放、手工初始化）通常是 with torch.no_grad(): 里原地改：

opt = torch.optim.Adam(model.parameters(), lr=0.1)
before = id(model.fc1.weight)
with torch.no_grad():
  model.fc1.weight.mul_(0.5)
print("same object after mul_:", id(model.fc1.weight) == before,
      "| still in optimizer:", any(p is model.fc1.weight for g in opt.param_groups for p in g["params"]))
new = nn.Parameter(model.fc1.weight * 0.5)
print("nn.Parameter(w * 0.5): is_leaf =", new.is_leaf, "requires_grad =", new.requires_grad)

same object after mul_: True | still in optimizer: True
nn.Parameter(w * 0.5): is_leaf = True requires_grad = True

还是同一个对象，优化器里的状态跟着它。新建一个 nn.Parameter 也没问题，它自己就是叶子。burn 没有原地改，只能换掉 Param，有两个入口，结果不一样：

let old = model.fc1.weight.id;
model.fc1.weight = model.fc1.weight.map(|w| w * 0.5);
println!("map:         id {old} -> {}", model.fc1.weight.id);
let old = model.fc1.weight.id;
model.fc1.weight = Param::from_tensor((model.fc1.weight.val() * 0.5).detach());
println!("from_tensor: id {old} -> {}", model.fc1.weight.id);
// 不 detach：w * 0.5 是图里的中间结果
let w = model.fc1.weight.val();
let r = std::panic::catch_unwind(|| Param::from_tensor(w * 0.5));
println!("from_tensor(w * 0.5) without detach: {}", if r.is_ok() { "ok" } else { "panicked" });

map:         id gjncvcpccv0p8 -> gjncvcpccv0p8
from_tensor: id gjncvcpccv0p8 -> ad7sarf0usrfo

thread 'main' (8622099) panicked at .../burn-autodiff-0.22.1/src/tensor.rs:190:17:
Can't convert a non leaf tensor into a tracked tensor
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
from_tensor(w * 0.5) without detach: panicked

Param::from_tensor 生成一个新的 ParamId。在优化器建好、跑过几步之后这样换，Adam 存的动量按旧 id 记着，新参数从零开始（优化器状态见优化器状态）。它还对传进来的张量调 require_grad()，在 autodiff 设备上，w * 0.5 是带历史的中间结果，于是 panic，信息和打开求导里对中间结果 require_grad() 的那条一样。nn.Parameter 会自己切断历史，from_tensor 不会，要先 detach()。
Param::map 保住 id，但不管新值是不是叶子。

map 不报错，所以更要看结果能不能训练（bin remap，三种写法各建一个新模型，前向、backward、GradientsParams::from_grads）：

let variants: [(&str, fn(Param<Tensor<2>>) -> Param<Tensor<2>>); 3] = [
  ("map(w * 0.5)", |p| p.map(|w| w * 0.5)),
  ("map((w * 0.5).detach())", |p| p.map(|w| (w * 0.5).detach())),
  ("map((w * 0.5).detach().require_grad())", |p| p.map(|w| (w * 0.5).detach().require_grad())),
];

map(w * 0.5)                             is_tracked=true  is_require_grad=false GradientsParams::len()=3
map((w * 0.5).detach())                  is_tracked=false is_require_grad=false GradientsParams::len()=3
map((w * 0.5).detach().require_grad())   is_tracked=true  is_require_grad=true  GradientsParams::len()=4

前两种写法，fc1.weight 都不再是要梯度的叶子，GradientsParams 只有 3 个，这个参数从此不会被更新，也不报错。第一种的梯度沿着 * 0.5 流回了旧的叶子，而旧叶子已经不在模型里（bin oldleaf：先留住旧的 weight.val()，map(w * 0.5) 之后 backward，旧句柄取到 Some([4, 8])，新参数取到 None）。只有第三种和 PyTorch 的 with torch.no_grad(): w.mul_(0.5) 等价：id 不变、值变了、仍是要梯度的叶子。burn 的优化器自己换参数时做的是同一件事：在内层后端上算出新值，ParameterContext::restore 用 Tensor::from_inner 抬回 autodiff，原来要梯度的再 require_grad()。

动脑筋答案

let bn = BatchNormConfig::new(8).init(&device);
let mut named = Named::default();
bn.visit(&mut named);
println!("BatchNorm(8) num_params = {}", bn.num_params());

gamma            [8]      require_grad=true  id=9ft6hq8vh04n8
beta             [8]      require_grad=true  id=8ddgm7n91esna
running_mean     [8]      require_grad=false id=u8lofguf05e7g
running_var      [8]      require_grad=false id=m2497vq8m9tfm
BatchNorm(8) num_params = 32

是 32，不是 16。RunningState 的 visit 把当前值包成一个临时的 Param，照样调 visitor.visit_float，所以访问者分不出它和 gamma，num_params 把运行统计也算进去了。拿 num_params() 和 PyTorch 的 sum(p.numel() for p in model.parameters()) 对数时，模型里每个 BatchNorm 会多出 2 × num_features。
区分的办法是 require_grad：运行统计永远是 false。GradientsParams 正是靠这一点跳过它们，优化器不会更新运行统计。
访问者也看不到 Dropout、BatchNorm 的 training 标志：它们走 visit_flag，Named 没有覆盖这个方法。

小结

named_parameters() 对一个十几行的 ModuleVisitor：路径名和 PyTorch 相同，形状转置，总数用 num_params()。
num_params() 把 BatchNorm 的 running_mean / running_var 也算进去（BatchNorm(8) 是 32，PyTorch 是 16）；require_grad=false 是区分它们的标记。
冻结：requires_grad_(False) 对 model.fc1 = model.fc1.no_grad() 或 set_require_grad_group(ParamGroup::from_predicate("fc1"), false)；冻结后训练一步，参数两边最大差 1.49e-8。冻结跨过 valid().train() 仍然有效。
改参数值用 param.map(|w| new.detach().require_grad())；少了这两步，参数静默退出训练；Param::from_tensor 换了 id，对中间结果直接 panic。