Rust in Everything
基于 burn 0.22.1

grad 返回 None

主线位置：求导主线的终点，backward() 返回 Gradients 之后按节点 id 取梯度。Gradients 怎么存、哪些条目会被删掉，见 backward 与 Gradients；哪些张量是叶子、Requirement 怎么区分，见 Autodiff 装饰器。

flowchart TB
  q1{"在 autodiff 设备上"} -->|否| none1["None"]
  q1 -->|是| q2{"is_require_grad"}
  q2 -->|否| none2["None：未追踪或中间结果"]
  q2 -->|是| q3{"这次反向走到它了吗"}
  q3 -->|"没有"| none3["None：不在容器里"]
  q3 -->|"走到了"| q4{"同一个 Gradients"}
  q4 -->|否| none4["None"]
  q4 -->|是| some["Some 梯度"]

实验 crate exp-burn/pb-pit，bin grads，Device::flex().autodiff()，debug，M2 Max、macOS 27、rustc 1.99.0。

现象

一个 loss 由五项相加，每一项让一个张量以不同的方式参与；w 是对照组：

let ad = Device::flex().autodiff();
let plain = Device::flex();

// 1. leaf on a device without autodiff
let x_plain = Tensor::<1>::from_floats([1.0, 2.0], &plain);
// 2. leaf on the autodiff device but no require_grad
let x_nograd = Tensor::<1>::from_floats([1.0, 2.0], &ad);
// 3. require_grad called after the forward used the tensor
let x_late = Tensor::<1>::from_floats([1.0, 2.0], &ad);
// the one that works
let w = Tensor::<1>::from_floats([3.0, 4.0], &ad).require_grad();
// 4. tracked leaf that does not reach the loss
let unused = Tensor::<1>::from_floats([5.0, 5.0], &ad).require_grad();
// 5. tracked leaf whose only path goes through detach()
let x_detached = Tensor::<1>::from_floats([1.0, 2.0], &ad).require_grad();
// 6. tracked leaf whose only path goes through an Int round trip
let x_int = Tensor::<1>::from_floats([1.0, 2.0], &ad).require_grad();

let x_detached_copy = x_detached.clone().detach();
let h = x_plain.clone() * w.clone();
let used_late = x_late.clone() * w.clone();
let x_late = x_late.require_grad();
let loss = (h.clone()
  + x_nograd.clone() * w.clone()
  + used_late
  + x_detached_copy.clone() * w.clone()
  + x_int.clone().int().float() * w.clone())
.sum();
let mut grads = loss.backward();
// ... 逐个打印 grad；然后 w.grad_remove(&mut grads)；
// 再对 (unused * 2).sum() 单独 backward 一次，得到 other

w (require_grad, on the path)      Some([5.0, 10.0])
1 x on Device::flex()              None
2 x without require_grad           None
3 require_grad after use           None
4 unused leaf                      None
5 path through detach()            None
  ...the detach() copy itself      Some([3.0, 4.0])
6 path through int().float()       None
7 intermediate h                   None
7 loss itself                      None
8 w after grad_remove              None
9 w, looked up in other Gradients  None
  unused, in other Gradients       Some([2.0, 2.0])
flags: x_nograd require_grad=false | h tracked=true require_grad=false

w 出现在五项里，每项的系数都是 [1, 2]，梯度 5 × [1, 2] = [5, 10]，计算本身没问题。其余每一行都是 None，返回值里看不出是哪一种。

出处

Tensor::grad 走 grad_impl：

fn grad_impl(p: &BridgeTensor, grads: &Gradients) -> Option<BridgeTensor> {
    // A non-float tensor — a packed base included — records no tape, so there
    // is no gradient to look up.
    let tensor = p.try_as_float()?;
    if tensor.autodiff == DispatchAutodiffContext::Disabled
        || !Dispatch::float_is_require_grad(tensor)
    {
        return None;
    }
    Dispatch::grad(tensor, grads.as_inner()).map(BridgeTensor::float)
}

前两个条件不满足直接返回 None；都满足才按节点 id 去 Gradients 里找，找不到也是 None。九种情形按落在哪一关分成三组：

#情形落在哪为什么
1张量建在 Device::flex() 上Disabled不在 autodiff 设备上，没有节点
2autodiff 设备上，没 require_grad()is_require_grad 为假Requirement::None，算子登记父节点时把它过滤掉
7中间结果 h、loss 本身is_require_grad 为假GradInBackward；它们的梯度在反向途中用完即删
3先参与了前向，之后才 require_grad()容器里没有前向时它还是 None，没有被登记成父节点
4叶子和 loss 无关容器里没有反向没走到它
5唯一的路径经过 detach()容器里没有detach 造了一个新 id 的叶子，梯度记在新 id 上
6唯一的路径经过 int().float()容器里没有Int 张量不进图，转回来的浮点张量是新的未追踪叶子
8已经 grad_remove 过容器里没有被取走了
9拿另一次 backward() 的 Gradients 查容器里没有每次 backward 一个新容器

几处值得对着源码看：

  • #3：算子输出的节点在 AutodiffTensor::from_parents 里建，父节点列表是 parent_nodes.iter().filter_map(|node| node.clone_if_require_grad())：前向那一刻不要梯度的输入根本不进父节点列表。之后的 require_grad 用同一个 id 新建一个 Grad 节点、登记一个根步骤，但早先的算子已经记好了，不会回头补。
  • #5：autodiff 后端的 float_detach 注释写着「we remove it from the graph, but we still want to keep the require_grad」：对一个要梯度的叶子 detach()，得到的是新 id 的、同样要梯度的叶子（上面输出里 the detach() copy itself 有梯度），原来那个句柄查不到。对中间结果 detach() 则得到不要梯度的叶子。
  • #7：中间结果要梯度，grad 帮不上；把它改写成叶子（h.detach().require_grad() 再参与后面的计算），代价是梯度不再往 h 之前传。

修法

按上面的分组倒着查，每组一个改法：

  • 第一组（Disabled、没 require_grad）：参数和输入建在 .autodiff() 设备上，要梯度的叶子在参与任何计算之前 require_grad()。模块里的 Param 由 Module::train() 或在 autodiff 设备上初始化来打开（见 Module 派生）。
  • 第二组（反向没走到）：从 loss 往回找这个张量，路上有没有 detach()、int()、argmax 一类整数结果、without_autodiff()，或者在另一个设备上算的分支。把要梯度的那一份留在图里，detach 之后的新句柄如果确实要用，就对新句柄取 grad。
  • 第三组（拿错了容器）：一次 backward 对应一个 Gradients；要合并两份损失的梯度，先把损失加起来再 backward（graph tape has already been consumed 里有这个写法），不要指望在两个容器之间来回查。

bin grads_fix 把第一组和 #7 的写法改过来：x 先 require_grad() 再参与前向；想要中间结果 h 的梯度，就把它截成一个新的叶子：

let w = Tensor::<1>::from_floats([3.0, 4.0], &ad).require_grad();
// require_grad before the tensor takes part in anything
let x = Tensor::<1>::from_floats([1.0, 2.0], &ad).require_grad();
// an intermediate whose gradient we want: make it a leaf of the rest
let h = (x.clone() * w.clone()).detach().require_grad();
let loss = (h.clone() * 2.0 + x.clone() * w.clone()).sum();
let grads = loss.backward();

w                        Some([1.0, 2.0])
x                        Some([3.0, 4.0])
h (detached leaf)        Some([2.0, 2.0])

三个都有梯度。代价也看得见：w 的梯度只剩 x * w 那一项的 [1, 2]，经过 h 的那一路被 detach 截断了。只想看一眼中间结果的梯度时这样做；训练时不要截。

怎么确认没中招

  • 训练循环里 GradientsParams::from_grads(grads, &model).len() 应当等于要训练的参数个数。是 0 或偏少，先打印每个参数的 is_require_grad() 和设备。
  • 对单个张量，按顺序问三件事：t.device().is_autodiff()、t.is_require_grad()、它到 loss 的路径上有没有切断梯度的操作。前两问为假就是第一组。
  • 模型换过设备之后梯度全为 0，是另一种情形：参数变成了 to_device 算子的输出，见 to_device 之后的模型。

小结

  • grad 返回 None 有九种情形，分三组：张量不是被追踪的叶子、反向没走到它、拿错了 Gradients。
  • 算子在前向时就决定了父节点列表，之后才 require_grad() 来不及。
  • detach() 一个要梯度的叶子，得到新 id 的叶子，梯度记在新句柄上。
  • 训练时用 GradientsParams::len() 一行就能发现「没有梯度」。

💬 关联讨论

发起讨论
加载中...