Rust in Everything
基于 burn 0.22.1

不求导的区域

主线位置：部署主线的「推理不带 autodiff」一步；也包括训练主线里不该进图的计算（指标、日志、冻结的特征）。
版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）
flowchart TB
  q{"要不要记图"} -->|"PyTorch"| ctx["with torch.no_grad"]
  q -->|"burn 推理"| plain["不带 autodiff 的设备"]
  q -->|"burn 已有模型"| valid["Module valid"]
  q -->|"burn 单个张量"| wa["without_autodiff"]
  q -->|"切断一段历史"| det["detach"]
  det -.->|"叶子仍要梯度"| warn["burn 照样记图"]

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

没有上下文，只有设备

PyTorch 用一个线程级的模式开关决定记不记图：with torch.no_grad(): 里的运算不记，出来以后照常。burn 没有这样的作用域，记不记图由每个张量自己带的 autodiff 标记决定，标记来自它所在的设备（见 打开求导）。所以 PyTorch 的几种写法在 burn 里换成「把张量或模型换到不带 autodiff 的那一边」：

w = torch.tensor([3.0, 4.0], requires_grad=True)
with torch.no_grad():
  a = w * 2
with torch.inference_mode():
  b = w * 2
d = w.detach()
print(f"no_grad:        requires_grad={a.requires_grad} grad_fn={a.grad_fn}")
print(f"inference_mode: requires_grad={b.requires_grad} is_inference={b.is_inference()}")
print(f"w.detach():     requires_grad={d.requires_grad} same storage={d.data_ptr() == w.data_ptr()}")

let w = Tensor::<1>::from_floats([3.0, 4.0], &ad).require_grad();
flags("w", &w);
flags("w.detach()", &w.clone().detach());
flags("(w*2).detach()", &(w.clone() * 2.0).detach());
flags("w.without_autodiff()", &w.clone().without_autodiff());
flags("w.inner()", &w.clone().inner());

no_grad:        requires_grad=False grad_fn=None
inference_mode: requires_grad=False is_inference=True
w.detach():     requires_grad=False same storage=True

w          is_autodiff=true  is_tracked=true  is_require_grad=true
w.detach() is_autodiff=true  is_tracked=true  is_require_grad=true
(w*2).detach() is_autodiff=true  is_tracked=false is_require_grad=false
w.without_autodiff() is_autodiff=false is_tracked=false is_require_grad=false
w.inner()  is_autodiff=false is_tracked=false is_require_grad=false

without_autodiff 把张量交给内层后端，去掉图的引用；之后只和普通张量一起算时，不经过 autodiff 那一层。inner 是同一个函数的另一个名字。去掉标记的张量再和 autodiff 张量一起运算，会被当成常量。

PyTorchburn
torch.inference_mode() 包住整个推理模型和输入建在 Device::flex() 这类不带 autodiff 的设备上
with torch.no_grad(): model(x)model.valid().forward(x.without_autodiff())
with torch.no_grad(): 里算指标对要记的张量 without_autodiff()
t.detach()中间结果 detach()；要梯度的叶子见下一节
📝NOTE
冷知识：burn 有过一个叫 no_grad 的张量方法

PR #5319（「Keep metric items on their device instead of syncing them to the host」，首次发布于 v0.22.0-pre.2）给 Tensor 加了 pub fn no_grad(self) -> Self，文档说它把张量从 autodiff 后端拿下来，「unlike inner, which panics on a tensor with no autodiff wrapper, this is safe to call anywhere」，用在 batcher 和指标管道里。名字和 PyTorch 的一样，做的却是现在 without_autodiff 做的事，不是一个作用域。四周后的 PR #5557（「feat(tensor): replace no_grad with explicit autodiff conversions」，首次发布于 v0.22.0-pre.4）把它改名为 without_autodiff，inner() 改成直接调它，同时新增了 Tensor::autodiff()；burn-book 里的「you can call inner() to obtain the inner tensor」改成「without_autodiff() ... The historical inner() method is equivalent.」。同一个提交还删掉了文件末尾一条从 0.21 就在的 TODO：「so it might be cleaner and easier to manage the device etc. if we provide a method to scope the autodiff?」它设想的是给 autodiff 划一个作用域，接近 torch.no_grad() 的形态，到 0.22.1 也没有这样的 API。Tensor::no_grad 只在 pre.2 和 pre.3 两个预发布版本里存在过；Module::no_grad 是另一回事，见下文。

detach 不一样

上面的输出里，PyTorch 的 w.detach() 不要梯度，burn 的 w.detach() 仍然 is_require_grad=true。burn 的 float_detach 是这样写的：

    fn float_detach(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        // When we detach a tensor, we remove it from the graph, but we still want to keep the
        // `require_grad` (and `distributed`) setting.
        let is_require_grad = Self::float_is_require_grad(&tensor);

        let distributed_params = tensor.node.distributed_params.clone();

        let mut tensor = AutodiffTensor::new(tensor.primitive);

        if is_require_grad {
            tensor = tensor.require_grad();
        }
        // ...
        tensor
    }

两边的 detach 都是「从这里切断历史」。区别在叶子：PyTorch 的结果一律不要梯度；burn 对一个要梯度的叶子 detach()，得到一个新 id、同样要梯度的叶子，用它做的运算照样记图，梯度记在新句柄上（grad 返回 None 的 #5）。对中间结果 detach()，两边一样，得到不追踪的张量。所以 PyTorch 里「param.detach() 拿来做一次不求导的计算」的写法，搬到 burn 要写成 param.without_autodiff()。

一条不要梯度的支路，两边写法对应、结果相同：

x = torch.tensor([1.0, 2.0], requires_grad=True)
with torch.no_grad():
  side = x * 10
loss = (side * w).sum() + (x * w).sum()
loss.backward()
print("x.grad =", x.grad, "(only the x*w path)")

let x = Tensor::<1>::from_floats([1.0, 2.0], &ad).require_grad();
let side = (x.clone() * 10.0).detach();
let grads = ((side * w.clone()).sum() + (x.clone() * w.clone()).sum()).backward();
println!("x.grad = {:?} (only the x*w path)", v(x.grad(&grads)));

x.grad = tensor([3., 4.]) (only the x*w path)

x.grad = Some([3.0, 4.0]) (only the x*w path)

推理：Module::valid

PyTorch 推理通常是 model.eval() 加 torch.no_grad()。burn 里两件事都由 Module::valid 做：它返回一份去掉 autodiff 的模型，同时关掉 dropout 一类的训练行为（训练行为的部分见训练和推理模式）：

let fc = LinearConfig::new(2, 2).init(&ad);
let input = Tensor::<2>::from_floats([[1.0, 2.0]], &ad);
flags("fc(x)", &fc.forward(input.clone()));
let fc_valid = fc.valid();
flags("valid fc.w", &fc_valid.weight.val());
flags("valid fc(x)", &fc_valid.forward(input.clone().without_autodiff()));
flags("valid fc(ad x)", &fc_valid.forward(input));
flags("train fc.w", &fc_valid.train().weight.val());

fc(x)      is_autodiff=true  is_tracked=true  is_require_grad=false
valid fc.w is_autodiff=false is_tracked=false is_require_grad=false
valid fc(x) is_autodiff=false is_tracked=false is_require_grad=false
valid fc(ad x) is_autodiff=true  is_tracked=false is_require_grad=false
train fc.w is_autodiff=true  is_tracked=true  is_require_grad=true
  • valid() 之后参数不在 autodiff 上，输入也 without_autodiff()，输出就是普通张量，和 PyTorch no_grad 下的输出一样。
  • 输入还在 autodiff 设备上时，输出仍带 autodiff 标记，但不追踪：参数这时是常量，没有要梯度的东西。图上不会留下东西，只是每个算子多走一层 autodiff 的分派。
  • valid() 借用 &self、返回新模型，train() 再把梯度设置打开。

Module::no_grad 名字像 torch.no_grad，作用却是 PyTorch 的 requires_grad_(False)：把所有参数设成不要梯度，dropout 照常，文档建议推理用 valid，no_grad 用于部分微调（见参数与遍历）。

只前向不反向：图一直占着内存

PyTorch 的图挂在输出张量的 grad_fn 链上，输出没人引用，图就随之释放。burn 的图存在进程级的结构里，只在 backward 结束时清理（机制见 burn 卷 backward 与 Gradients 的「什么时候释放」）。在 autodiff 设备上只做推理，差别就出来了。

实验：一个 1024×1024 的 f32 参数（4 MiB）要梯度，前向是 ((p.exp() * p) @ p).sum()，重复 20 次。每种情形单独起一个进程，先做一次前向加反向预热，再用 ps -o rss= 读常驻内存的增量：

p = torch.randn(n, n, requires_grad=True)
def forward(p):
  return ((p.exp() * p) @ p).sum()
forward(p).backward()  # warm up allocator and kernels
p.grad = None
base = rss_mib()
kept = []
for i in range(20):
  if mode == "drop":
    forward(p)
  elif mode == "keep":
    kept.append(forward(p))
  elif mode == "detach-keep":
    kept.append(forward(p.detach()))
  elif mode == "nograd-keep":
    with torch.no_grad():
      kept.append(forward(p))

let p = Tensor::<2>::random([1024, 1024], Distribution::Default, &ad).require_grad();
let _ = forward(&p).backward(); // warm up
let base = rss_mib();
let mut kept = Vec::new();
for _ in 0..20 {
  match mode.as_str() {
    "drop" => drop(forward(&p)),
    "keep" => kept.push(forward(&p)),
    "without-autodiff" => drop(forward(&p.clone().without_autodiff())),
    "detach" => drop(forward(&p.clone().detach())),
    _ => unreachable!("unknown mode"),
  }
}
// "drop" 之后再做一次 backward，再前向 20 次

drop         20 forwards: rss +   0.0 MiB
keep         20 forwards: rss + 144.6 MiB
detach-keep  20 forwards: rss +   0.0 MiB
nograd-keep  20 forwards: rss +   0.0 MiB

drop             20 forwards: rss + 147.0 MiB
drop             then one backward: rss + 163.1 MiB
drop             20 more forwards: rss + 163.1 MiB
keep             20 forwards: rss + 147.1 MiB
without-autodiff 20 forwards: rss +   0.1 MiB
detach           20 forwards: rss + 145.0 MiB
  • 每次前向为反向保存两个 4 MiB 的张量（exp 的输出给乘法用，乘积给矩阵乘用），20 次约 160 MiB；两边的 145～147 MiB 都在这个量级，比 160 少的部分是预热时释放、又被复用的内存。
  • PyTorch 只有把输出留在手里（keep）才涨；丢掉输出（drop）、detach、no_grad 都不涨。
  • burn 丢掉输出也涨（drop），和留着一样：保存的张量挂在图里，不挂在 loss 上。detach 也涨，原因是上一节说的，p.detach() 还是要梯度的叶子。只有 without_autodiff() 不涨。
  • 做一次 backward 以后，再前向 20 次 RSS 不再增长：清理把前 20 次的图释放了，内存被重新使用。RSS 本身没有降，是因为分配器没有把页还给系统，所以这里看的是「后面还涨不涨」。

推论：部署主线上，推理不要在 autodiff 设备上做。模型用 valid()（或者一开始就建在不带 autodiff 的设备上），输入用普通设备或 without_autodiff()；PyTorch 里「忘了 no_grad 只是慢一点」的习惯，在 burn 里会变成内存一直涨，直到下一次 backward。

小结

  • burn 没有 torch.no_grad() 式的作用域：记不记图看张量的 autodiff 标记，推理用不带 autodiff 的设备、Module::valid() 或 without_autodiff()。
  • detach 对中间结果两边一样；对要梯度的叶子，burn 保留 require_grad，结果照样记图。
  • Module::no_grad 对应的是 requires_grad_(False)，不是推理模式。
  • PyTorch 丢掉输出图就释放；burn 的图只在 backward 结束时清理，在 autodiff 设备上只前向，内存一直涨（实测 20 次 +147 MiB）。

💬 关联讨论

发起讨论
加载中...