Rust in Everything
基于 burn 0.22.1

所有权

PyTorch 的张量是 Python 对象,名字只是引用:c = a + b 之后 a 照常可用,b = a 让两个名字指向同一块内存,a.add_(1) 改的东西两个名字都看得见。burn 的 Tensor 是一个 Rust 值,大多数方法签名是 fn op(self, ...) -> Self。这一部分只讲由此带来的两处变化,Rust 所有权本身不展开。

  1. 运算消费张量:a + b 之后 a 已被移走,rustc 报 E0382;复用要 clone(),而 clone() 只加引用计数,64 MiB 的张量 clone 一次不分配数据缓冲,PyTorch 的 clone() 则多占 64 MiB。模块里 forward(&self) 用字段要 clone(),Param::val() 就是这么做的。
  2. 原地操作:PyTorch 的 add_ 改了 autograd 存下的张量,backward 时报版本计数错误;burn 没有公开的原地 API,x = x + 1 在缓冲只有自己持有时由后端原地写,autograd 存着它时就另开一块,梯度不受影响。
flowchart TB
  op["x 按值进入运算"] --> uniq{"缓冲只有一个持有者"}
  uniq -->|是| inplace["后端原地写 不分配"]
  uniq -->|否| fresh["另开缓冲写结果"]
  clone["clone 或 autodiff 存下"] -.->|计数加一| uniq

两页都以 Flex 设备为准。原地写在 burn 内部是怎么判断的,见 burn 卷 Tensor 与 Kind 的「按值消耗 self,clone 只加引用计数」。

💬 关联讨论

发起讨论
加载中...