Rust in Everything
基于 burn 0.22.1

backward 与梯度

主线位置：训练主线的 backward 一步：从 loss 算出梯度，交给优化器之前。
版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）
flowchart TB
  subgraph pt ["PyTorch"]
    pl["loss.backward"] -->|"累加"| pg["x.grad 属性"]
    pz["zero_grad"] -->|"清零"| pg
  end
  subgraph bn ["burn"]
    bl["loss.backward"] -->|"返回"| bg["Gradients 值"]
    bg --> bx["x.grad 加 grads"]
    bg --> gp["GradientsParams"]
  end

实验代码：PyTorch 侧 bg.py、accum.py，burn 侧 crate autodiff-exp 的 bin bg、accum（feature 同模板，debug）。平台 M2 Max、macOS 27、rustc 1.99.0。

梯度挂在哪

同一个前向 (x * x).sum()，重新算两遍、各 backward 一次：

x = torch.tensor([1.0, 2.0], requires_grad=True)
for i in range(2):
  (x * x).sum().backward()
  print(f"after backward {i}: x.grad = {x.grad}")

let ad = Device::flex().autodiff();
let x = Tensor::<1>::from_floats([1.0, 2.0], &ad).require_grad();
let mut all = Vec::new();
for i in 0..2 {
  let grads = (x.clone() * x.clone()).sum().backward();
  println!("backward {i}: x.grad = {:?}", v(x.grad(&grads)));
  all.push(grads);
}
let sum = x.grad(&all[0]).zip(x.grad(&all[1])).map(|(a, b)| a + b);
println!("summed by hand: {:?}", v(sum));

after backward 0: x.grad = tensor([2., 4.])
after backward 1: x.grad = tensor([4., 8.])

backward 0: x.grad = Some([2.0, 4.0])
backward 1: x.grad = Some([2.0, 4.0])
summed by hand: Some([4.0, 8.0])

PyTorch 的 backward() 返回 None，梯度写进叶子的 .grad 属性，已经有值就加上去，所以第二次是 [4, 8]。训练循环每步要 optimizer.zero_grad()，原因就在这里。

burn 的 Tensor::backward 返回一个 Gradients 值，叶子上什么也不写。每次 backward 都是一个新容器，两次互不相干，第二次还是 [2, 4]。要哪个张量的梯度，拿这个容器去问它：x.grad(&grads) 返回 Option<Tensor>，取的是一个句柄，梯度还留在容器里；只用一次时用 grad_remove 拿走。容器按节点 id 存、中间结果的梯度在反向途中就删掉，见 burn 卷的 backward 与 Gradients。

所以 burn 里没有 zero_grad：没有地方需要清零。上一步的 Gradients 在优化器用完以后被丢掉，下一步的 backward 再给一个新的。训练主线里这个容器接着交给 GradientsParams::from_grads(grads, &model)，按参数 id 重新组织后给优化器，见训练一步。

PyTorch 里和 burn 最像的是函数式的 torch.autograd.grad：它也把梯度作为返回值交出来，不碰 .grad：

h = x * x
(gx,) = torch.autograd.grad((h * 3).sum(), [x])
print("autograd.grad:", gx, "| x.grad =", x.grad)

autograd.grad: tensor([ 6., 12.]) | x.grad = None

主线：梯度累积

PyTorch 里梯度累积就是利用「不清零就累加」：把一批拆成几份，每份 backward 一次，中间不 zero_grad。burn 每次 backward 是新容器，累加要自己做，burn::optim 里有现成的 GradientsAccumulator。

用主线的两层 MLP（4→8→3，交叉熵，权重和 A1 相同），4 行的一批拆成两个 2 行的半批，每个半批的 loss 除以 2，累加出来应该等于整批一次 backward 的梯度：

ce = nn.functional.cross_entropy
model.zero_grad()
ce(model(x), y).backward()
full = grads()

# two half batches, no zero_grad in between: .grad accumulates
model.zero_grad()
for part in (slice(0, 2), slice(2, 4)):
  (ce(model(x[part]), y[part]) / 2).backward()
accum = grads()

let full = GradientsParams::from_grads(ce.forward(model.forward(x.clone()), y.clone()).backward(), &model);

// two half batches: each backward is a fresh container, sum them explicitly
let mut acc = GradientsAccumulator::<Mlp>::new();
for part in [0..2, 2..4] {
  let loss = ce.forward(model.forward(x.clone().slice(part.clone())), y.clone().slice(part)) / 2.0;
  acc.accumulate(&model, GradientsParams::from_grads(loss.backward(), &model));
}
let accum = acc.grads();

accumulate 遍历模型的每个参数，按 ParamId 从新的 GradientsParams 里取出梯度，和已有的相加；grads() 交出累加结果并把自己清空，相当于 PyTorch 那边 step() 之后的 zero_grad()。两边各输出 full 和 accum 两组梯度（权重转回 [out, in]），用 compare.py 对四组：

== full_pt vs full_burn
fc1.weight               n= 32  max|diff|=7.45e-09
fc1.bias                 n=  8  max|diff|=7.45e-09
fc2.weight               n= 24  max|diff|=1.49e-08
fc2.bias                 n=  3  max|diff|=1.49e-08
worst=1.49e-08 tol=1e-06 OK
== acc_pt vs acc_burn
fc1.weight               n= 32  max|diff|=7.45e-09
fc1.bias                 n=  8  max|diff|=7.45e-09
fc2.weight               n= 24  max|diff|=1.49e-08
fc2.bias                 n=  3  max|diff|=1.49e-08
worst=1.49e-08 tol=1e-06 OK
== full_pt vs acc_pt
fc1.weight               n= 32  max|diff|=7.45e-09
fc1.bias                 n=  8  max|diff|=0.00e+00
fc2.weight               n= 24  max|diff|=1.86e-09
fc2.bias                 n=  3  max|diff|=7.45e-09
worst=7.45e-09 tol=1e-06 OK
== full_burn vs acc_burn
fc1.bias                 n=  8  max|diff|=0.00e+00
fc1.weight               n= 32  max|diff|=7.45e-09
fc2.bias                 n=  3  max|diff|=7.45e-09
fc2.weight               n= 24  max|diff|=7.45e-09
worst=7.45e-09 tol=1e-06 OK
  • 整批梯度两边最大差 1.49e-8，和 A1 的结果一样；累积出来的梯度两边也是 1.49e-8。容差 1e-6 通过。
  • 同一边「整批」和「两个半批相加」也不是逐位相同（最大 7.45e-9）：整批时交叉熵对 4 行求平均、矩阵乘按 4 行求和，拆开后变成两次 2 行的求和再相加，f32 加法的顺序变了。两边不为 0 也是同一个原因：两个框架的矩阵乘和归约内核求和顺序不同。
  • 这些梯度的绝对值都在 0.2 以下，逐个参数组看，整批两边的最大差是 1 到 4 个 ulp（f32 在该数值上的最低位）。

非标量输出

PyTorch 的 backward() 只接受标量，非标量要显式给出 gradient；burn 对任何形状都直接 backward，种子是同形状的全 1 张量（Gradients::new），等于先 sum() 再求导：

y = x * x
try:
  y.backward()
except RuntimeError as e:
  print("non-scalar backward:", e)
y.backward(torch.ones_like(y))
print("y.backward(ones): x.grad =", x.grad)
x.grad = None
y = x * x
y.backward(torch.tensor([1.0, 10.0]))
print("y.backward([1, 10]): x.grad =", x.grad)

let y = x.clone() * x.clone();
let grads = y.backward();
println!("y.backward() on [2]: x.grad = {:?}", v(x.grad(&grads)));
let y = x.clone() * x.clone();
let weights = Tensor::<1>::from_floats([1.0, 10.0], &ad);
let grads = (y * weights).sum().backward();
println!("(y * [1, 10]).sum(): x.grad = {:?}", v(x.grad(&grads)));

non-scalar backward: grad can be implicitly created only for scalar outputs
y.backward(ones): x.grad = tensor([2., 4.])
y.backward([1, 10]): x.grad = tensor([ 2., 40.])

y.backward() on [2]: x.grad = Some([2.0, 4.0])
(y * [1, 10]).sum(): x.grad = Some([2.0, 40.0])

burn 的 backward 不接受外来的梯度，PyTorch 的 y.backward(v)（向量-雅可比积）在 burn 里写成 (y * v).sum().backward()，v 不要梯度，结果相同。反过来，从 PyTorch 搬代码时如果自己写的 loss 忘了最后的 mean()，PyTorch 会报上面那个错，burn 不会：它按求和处理，梯度是按平均算的 batch 大小倍。

中间结果的梯度

h = x * x
loss = (h * 3).sum()
loss.backward()
print("h.grad =", h.grad)

let h = x.clone() * x.clone();
let grads = (h.clone() * 3.0).sum().backward();
println!("h.grad = {:?} | x.grad = {:?}", v(h.grad(&grads)), v(x.grad(&grads)));

.../bg.py:30: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more information. (Triggered internally at ...)
  print("h.grad =", h.grad)
h.grad = None

h.grad = None | x.grad = Some([6.0, 12.0])

两边都是 None。PyTorch 给一条警告并提示 retain_grad()；burn 不警告，也没有 retain_grad：中间结果的梯度在反向途中被用掉就删了。grad 返回 None 的全部情形和判断顺序见 grad 返回 None，那里也有「把中间结果 detach().require_grad() 成新叶子」的写法和它的代价。

二阶导数

PyTorch 用 create_graph=True 让梯度本身也带图，可以再求一次导。burn 的梯度是不带 autodiff 的普通张量：

(g,) = torch.autograd.grad((x ** 3).sum(), [x], create_graph=True)
print("dy/dx =", g, "requires_grad =", g.requires_grad)
(g2,) = torch.autograd.grad(g.sum(), [x])
print("d2y/dx2 =", g2)

let grads = (x.clone().powi_scalar(3)).sum().backward();
let g = x.grad(&grads);
if let Some(g) = &g {
  flags("x.grad", g);
}
println!("dy/dx = {:?}", v(g.clone()));
attempt("backward on the gradient", || g.map(|g| g.sum().backward()));

dy/dx = tensor([ 3., 12.], grad_fn=<MulBackward0>) requires_grad = True
d2y/dx2 = tensor([ 6., 12.])

x.grad     is_autodiff=false is_tracked=false is_require_grad=false
dy/dx = Some([3.0, 12.0])

thread 'main' (8533725) panicked at .../burn-tensor-0.22.1/src/tensor/api/autodiff.rs:75:9:
Tensor::backward requires a tracked autodiff tensor; call Tensor::autodiff().require_grad() on the source leaf before computing the output
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
backward on the gradient: panicked

grad 的文档写着「The returned gradient doesn't have an autodiff association.」。burn 的自动微分只做一阶（见 burn 卷自动微分的开头），没有 create_graph 的对应；Hessian 向量积、梯度惩罚（gradient penalty）这类要对梯度再求导的写法，0.22.1 搬不过来。

改梯度

PyTorch 里改梯度是原地改 .grad（x.grad.clamp_(-5, 5)，或 clip_grad_norm_）。burn 的梯度在容器里，取出来算好再用 grad_replace 放回去：

let mut grads = (x.clone() * x.clone() * 10.0).sum().backward();
let clipped = x.grad(&grads).map(|g| g.clamp(-5.0, 5.0));
if let Some(c) = clipped {
  x.grad_replace(&mut grads, c);
}
println!("after grad_replace(clamp): {:?}", v(x.grad(&grads)));
println!("grad_remove: {:?} then grad: {:?}", v(x.grad_remove(&mut grads)), v(x.grad(&grads)));

after grad_replace(clamp): Some([5.0, 5.0])
grad_remove: Some([5.0, 5.0]) then grad: None

原来的梯度是 20x = [20, 40]，截到 [5, 5]。训练时按参数做梯度裁剪不用自己写：burn 的优化器配置（SgdConfig、AdamConfig 等）有 gradient_clipping 字段，类型是 GradientClippingConfig。

小结

  • PyTorch 把梯度累加进 .grad，每步要 zero_grad；burn 的 backward() 返回新的 Gradients，按张量 x.grad(&grads) 取，不累加，没有 zero_grad。
  • 梯度累积用 GradientsAccumulator；主线 MLP 上整批和两个半批累积的梯度，两边最大差都是 1.49e-8（容差 1e-6）。
  • 非标量输出：PyTorch 报错，burn 当成求和；y.backward(v) 写成 (y * v).sum().backward()。
  • 中间结果没有梯度，burn 也没有 retain_grad；梯度不带图，没有二阶导数。
  • 改梯度用 grad 取出、grad_replace 放回。

💬 关联讨论

发起讨论
加载中...