PyTorch 叶子张量

叶子张量是 PyTorch 计算图中的一个重要概念，叶子张量指的就是我们的模型参数，而模型参数一般都是我们自己创建的 requires_grad=True 的张量。它位于整个计算图的开始位置，比如下面这个例子：

import torch

def test01():

    a = torch.tensor([[2.0, 1.0]], requires_grad=True)
    b = torch.tensor([[1.0, 1.0]], requires_grad=True)
    c = a * 3
    d = b + 4
    e = c + d
    f = torch.mean(e)

    print(a.is_leaf)
    print(b.is_leaf)
    print(c.is_leaf)
    print(d.is_leaf)
    print(e.is_leaf)
    print(f.is_leaf)

if __name__ == '__main__':
    test01()

程序输出结果：

True
True
False
False
False
False

我们发现，除了 a、b 是叶子张量，其他张量都不是。从这里，我们还可以看到一个规律：非叶子张量都是通过其他张量计算得来的，而叶子张量是自己创建的。

叶子张量概念可以帮助 PyTorch 确定哪些内存可以释放。看下下面这个熟悉的报错，大致意思是第二次反向计算梯度时报错，报错的原因则是：反向传播是基于链式求导，如果中间梯度被释放了（这里指的是 c、d、e、f 的梯度），就无法再计算叶子张量的梯度。

def test02():

    a = torch.tensor([[2.0, 1.0]], requires_grad=True)
    b = torch.tensor([[1.0, 1.0]], requires_grad=True)
    c = a * 3
    d = b + 4
    e = c + d
    f = torch.mean(e)

    # 第一次反向传播
    f.backward()
    # tensor([[1.5000, 1.5000]]) tensor([[0.5000, 0.5000]]) None None None None
    print(a.grad, b.grad, c.grad, d.grad, e.grad, f.grad)

我们发现 backward 之后，除了叶子张量，其他张量的梯度就自动设置为 None。如果我们想要保留中间梯度值，可以使用对非叶子张量使用 retain_grad 方法：

def test02():

    a = torch.tensor([[2.0, 1.0]], requires_grad=True)
    b = torch.tensor([[1.0, 1.0]], requires_grad=True)
    c = a * 3
    d = b + 4
    e = c + d
    f = torch.mean(e)

    # 保留非叶子张量梯度
    c.retain_grad()
    d.retain_grad()
    e.retain_grad()
    f.retain_grad()

    # 第一次反向传播
    f.backward()
    # tensor([[1.5000, 1.5000]])
    # tensor([[0.5000, 0.5000]])
    # tensor([[0.5000, 0.5000]])
    # tensor([[0.5000, 0.5000]])
    # tensor([[0.5000, 0.5000]])
    # tensor(1.)
    print(a.grad, b.grad, c.grad, d.grad, e.grad, f.grad)

这时，你会发现反向传播完也不会清空中间变量梯度了。但是这也会导致更多的内存占用。再看下面的报错，应该就清楚原因是什么，就是因为默认情况下，第一次 backward 清空了中间梯度，导致第二次调用 backward 函数时候，不存在中间梯度，导致无法完成反向梯度计算。

RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.

从上面的过程，也可以看出，计算图的中间过程产生的张量尽量不要去修改，否则可能导致反向梯度计算时出现问题。事实上，PyTorch 也会跟踪这些张量的变化。比如：

def test03():

    a = torch.tensor([[2.0, 1.0]], requires_grad=True)
    b = torch.tensor([[1.0, 1.0]], requires_grad=True)
    c = a * 3
    d = b + 4
    e = c + d
    f = torch.mean(e)

    # 在反向计算梯度前，修改了张量 a 的值
    # 这是一种 in place 的操作，会修改张量本身的值
    a[0] = 100

    f.backward()

程序执行结果：

RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.

这个问题可以理解为，计算图已经构建好，准备要反向计算梯度了，结果你把需要的一些张量的值修改了。这可能导致反向计算的梯度错误，所以，PyTorch 不允许这些 in-place 操作。

对于计算图中的中间张量，也是不允许修改，如下所示：

def test04():

    a = torch.tensor([[2.0, 1.0]], requires_grad=True)
    b = a ** 3
    c = b ** 2
    d = torch.mean(c)

    # 中间结果被修改
    b[0][0] = 100.0

    d.backward()

报错内容如下，大致意思是：用于计算梯度的张量被使用了 in-place 方式修改了。

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [1, 2]], which is output 0 of torch::autograd::CopySlices, is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

最后稍微总结下：

叶子张量一般指的就是我们的模型参数，它需要计算梯度，存储梯度，而非叶子张量的梯度计算是为了计算叶子张量的梯度（中间结果），所以为了节省内存，非叶子张量的梯度在反向传播之后会被清空。如果此时再次进行反向梯度计算则会报错。
在反向传播之前，不允许对计算图中的叶子张量和非叶子张量进行 in place 修改值的操作，这会导致反向梯度计算时出错。这里需要注意一点，有时你会发现你修改了中间张量，但是并不会报错，这是因为，这些张量不影响到梯度的计算，PyTorch 就没有检查该张量是否被修改。

PyTorch 叶子张量

取消回复