考虑到Rust不允许多个可变引用，类似PyTorch的自动区分如何在Rust中工作

发布于10月07日

我主要是一个局外人，试图了解铁 rust 是否适合我的项目.

在Rust中有一些框架可以进行自动区分.具体来说，candle和其他一些项目，我认为，根据他们的描述，不知何故以一种类似于 torch 的方式做到了这一点.

然而，我知道Rust不允许多个可变引用.而这似乎就是类似于PyTorch的自动区分所需要的:

x = torch.rand(10) # an array of 10 elements
x.requires_grad = True

y = x.sin()
z = x**2

104、y和z必须保持对x的可变引用，因为您可能希望反向传播它们，这将修改x.grad.例如:

(y.dot(z)).backwards()
print(x.grad) # .backwards() adds a new field (an array) to x, without modifying it otherwise

既然Rust不允许多个可变引用，那么如何在Rust中实现类似的行为呢？

推荐答案

您是正确的，Ruust编译器强制一次只能有一个对值的可变引用，但有一个转义线:interior mutability模式.

该模式允许程序员构造数据 struct ，其规则在run time而不是编译时被判断.

标准库提供了许多实现内部可变性的容器，具有适合不同场景的不同使用模式.主要的例子有:

RefCell<T>，它允许单线程使用的运行时borrow 判断
RwLock<T>，它允许运行时borrow 判断多线程使用
Mutex<T>，一次只允许一个对其内容的引用

还有其他的--请参阅cell和sync的模块级别文档.

这一点如何适用于蜡烛？让我们来看看peek under the hood分:

pub struct Tensor_ {
    ...
    storage: Arc<RwLock<Storage>>,
    ...

支持张量的存储器的内容由RwLock保护.事实上，代码中紧靠其上的一些注释描述了 Select 此特定解决方案的原因--值得一读.

不仅如此，它还被包装在Arc<T>中-这意味着它实际上是一个分配的堆，引用计数值.该值可以有多个"所有者"，只有当最后一个所有者超出范围时，才会释放该值.

在反向传播的情况下，这是如何使用的？Tensor的backward()方法并不直接修改张量，而是返回包含计算出的梯度的GradStore.A GradStore又可能被A Optimizer消耗.Optimizer是一个特征，有两个不同的实现，所以让我们来看看SGD优化器:

    fn step(&mut self, grads: &candle::backprop::GradStore) -> Result<()> {
        for var in self.vars.iter() {
            if let Some(grad) = grads.get(var) {
                var.set(&var.sub(&(grad * self.learning_rate)?)?)?;
            }
        }
        Ok(())
    }

好的，这里的渐变被应用于大约Var个实例--这些是什么(定义为here)？

pub struct Var(Tensor);

好的，Tensor的包装.而set method是如何做好它的工作的？这一行是关键:

let (mut dst, layout) = self.storage_mut_and_layout();

这为我们提供了一个可变变量，它似乎代表了SET操作的目的地.这个storage_mut_and_layout()方法是做什么的？

let storage = self.storage.write().unwrap();

啊哈！它对上面看到的RwLock调用write()方法，存储位于其中.此方法的文档说明如下:

使用独占写入访问锁定此卢旺达Lock，从而阻止当前线程，直到它可以被获取.

当其他编写器或其他读取器目前有权访问该锁.

因此，总而言之:

backward()方法本身似乎不修改输入Tensor，但它返回包含渐变的数据 struct
使用Optimizer将渐变应用到Tensor.
Optimizer使用set方法来改变Tensor,Tensor在幕后使用保护它的RwLock上的write()方法获得对Tensor‘S数据存储的可变访问.