所有权
PyTorch 的张量是 Python 对象,名字只是引用:c = a + b 之后 a 照常可用,b = a 让两个名字指向同一块内存,a.add_(1) 改的东西两个名字都看得见。burn 的 Tensor 是一个 Rust 值,大多数方法签名是 fn op(self, ...) -> Self。这一部分只讲由此带来的两处变化,Rust 所有权本身不展开。
- 运算消费张量:
a + b之后a已被移走,rustc 报 E0382;复用要clone(),而clone()只加引用计数,64 MiB 的张量 clone 一次不分配数据缓冲,PyTorch 的clone()则多占 64 MiB。模块里forward(&self)用字段要clone(),Param::val()就是这么做的。 - 原地操作: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 只加引用计数」。
💬 关联讨论
发起讨论加载中...