第二次 backward 后梯度翻倍
主线位置:eager 主线的最后两步,
loss.backward(&mut engine)→x.grad()。梯度槽、trace 和函数式grad的机制见 eager AD。
flowchart TB
step["每一步 新的 loss"] --> bw["backward"]
bw --> slot{"槽里已有梯度"}
slot -->|没有| put["存入"]
slot -->|有| add["相加"]
add --> grow["n 步之和"]
clear["clear_grad"] --> slot
实验 crate exp-tenferro8/pit2-exp,bin p18_accumulate(参数 clear 是修好的版本),CPU(CpuEngine::serial()),debug,M2 Max、macOS 27、rustc 1.99.0。
现象
一个最小的「训练循环」:w 是变量,x 是常数,每一步重新算 loss = sum(w * x) 再 backward。梯度应当每步都是 x:
let rt = EagerRuntime::new()?;
let mut engine = CpuEngine::serial();
let w = rt.variable_from(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0])?)?;
let x = rt.constant_from(Tensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0])?)?;
for step in 1..=3 {
let (h, loss) = rt.with_eager_session(&mut engine, |s| {
let h = s.mul(&w, &x)?;
let loss = s.reduce_sum(&h, None)?;
Ok::<_, tenferro_ad::Error>((h, loss))
})?;
loss.backward(&mut engine)?;
println!("step {step}: w.grad = {:?} h.grad = {:?} x.grad = {:?}", g(&w)?, g(&h)?, g(&x)?);
}
step 1: w.grad = Some([10.0, 20.0, 30.0]) h.grad = None x.grad = None
step 2: w.grad = Some([20.0, 40.0, 60.0]) h.grad = None x.grad = None
step 3: w.grad = Some([30.0, 60.0, 90.0]) h.grad = None x.grad = None
每一步的 loss 都是新建的,w.grad() 却是前几步之和。拿它更新参数,第 n 步的步长放大 n 倍。没有报错、没有警告。
同一个程序末尾,几个相邻的误会:
h.backward(): Err(grad requires a scalar output, got shape [3])
NonScalarGrad { shape: [3] }
rt.grad(loss, h): Err(grad (GraphBuild): invalid argument wrt: grad is inactive for Input(User { id: 16 }); use grad_optional to observe the inactive state)
rt.grad_optional(loss, h): is_some = false
rt.grad(loss, w) = [10.0, 20.0, 30.0], w.grad still Some([30.0, 60.0, 90.0])
- 中间结果
h和常数x没有梯度槽,grad()是None,不是零。 h的形状是[3],backward只收标量,报NonScalarGrad;非标量要用backward_with(&mut engine, &cotangent)给出种子。- 函数式的
rt.grad(&mut engine, &loss, &h)对中间结果求导报「grad is inactive」,消息里是h的内部键Input(User { id: 16 }),看不出指的是哪个张量。 rt.grad(&mut engine, &loss, &w)返回一个新张量,不碰槽位:槽里还是三步之和。
出处
EagerTensor::backward 先检查标量,再从全 1 的种子做一次反向,最后由 store_grads 写进槽位。槽里已有值时在 engine 上相加:
let next = match current.as_ref() {
Some(existing) => {
let existing_read = existing.tensor_read("EagerRuntime::store_grads")?;
let incoming_read = TensorRead::from_tensor(incoming);
let tensor = session
.add_read(existing_read, incoming_read)
.map_err(Error::from)?;
AdValueRecord::from_tensor(tensor, "EagerRuntime::store_grads")?
}
这是有意的:backward 的文档示例就是在循环里调两次、断言梯度是 [4.0, 4.0, 4.0](一次是 2)。和 PyTorch 一样,累加让一个大 batch 可以拆成几次 backward 再一起更新。只有 variable_from 建的叶子有槽;中间结果和 constant_from 的常数没有,所以是 None。
函数式的 EagerRuntime::grad 拿到 None(wrt 不是活动的叶子)时,调 traced 一侧的 inactive_wrt_error 拼出这条 Validation 错误,grad_optional 则把同一个状态当成 Ok(None)。0.7.1 报的是 Internal("grad output is inactive ..."),提交 98319e01(「classify an inactive wrt as typed validation, not Internal」)把它改成了用户可以处理的校验错误。
修法
每一步 backward 之前清掉槽:
for step in 1..=3 {
w.clear_grad()?;
// ... 同上
}
step 1: w.grad = Some([10.0, 20.0, 30.0]) h.grad = None x.grad = None
step 2: w.grad = Some([10.0, 20.0, 30.0]) h.grad = None x.grad = None
step 3: w.grad = Some([10.0, 20.0, 30.0]) h.grad = None x.grad = None
- 一个变量用
clear_grad;整个运行时的所有槽用EagerRuntime::clear_grads,相当于 PyTorch 的optimizer.zero_grad()。 - 不想管槽位,就别用
backward:let dw = rt.grad(&mut engine, &loss, &w)?每次返回新的梯度张量(上面最后一行)。 - 要中间结果的梯度,用
rt.grad_optional先看它是不是活动的;需要真的梯度,就把它改成一个叶子再算后半段。
怎么确认没中招
- 训练循环里每一次
backward之前,搜一下有没有clear_grad/clear_grads。 - 第一步和第二步用同样的数据各跑一次,
grad()应当相等;第二步是第一步的两倍就是忘了清。 grad()返回None时先看这个张量是不是variable_from建的(tracks_grad()为真的不一定是叶子:中间结果h的tracks_grad()也是真)。
小结
- eager 的
backward把梯度加进叶子的槽位,不清就逐步累加;第 n 步的grad()是 n 步之和。 - 每步之前
clear_grad()(或运行时的clear_grads()),或者改用不碰槽位的rt.grad。 - 只有
variable_from的叶子有梯度;非标量要backward_with,对中间结果用rt.grad报grad is inactive ... use grad_optional。