评价此页

★ ★ ★ ★ ★

beginner/understanding_leaf_vs_nonleaf_tutorial

在 Google Colab 中运行

注意

跳转至页面底部下载完整示例代码。

理解 requires_grad, retain_grad, Leaf 和 Non-leaf 张量#

作者: Justin Silver

本教程通过一个简单的示例解释了 requires_grad、retain_grad、叶子张量 (leaf tensors) 和非叶子张量 (non-leaf tensors) 的细微差别。

开始之前，请确保您了解张量及其操作方法。对自动求导 (autograd) 工作原理的基础知识也会有所帮助。

设置#

首先，请确保已安装 PyTorch，然后导入必要的库。

import torch
import torch.nn.functional as F

接下来，我们实例化一个简单的网络以专注于梯度计算。该网络包含一个仿射层 (affine layer)，后接一个 ReLU 激活函数，最后是预测张量与标签张量之间的均方误差 (MSE) 损失。

\[\mathbf{y}_{\text{pred}} = \text{ReLU}(\mathbf{x} \mathbf{W} + \mathbf{b})\]

\[L = \text{MSE}(\mathbf{y}_{\text{pred}}, \mathbf{y})\]

请注意，参数（W 和 b）必须设置 requires_grad=True，这样 PyTorch 才能跟踪涉及这些张量的操作。我们将在未来的章节中详细讨论这一点。

# tensor setup
x = torch.ones(1, 3)                      # input with shape: (1, 3)
W = torch.ones(3, 2, requires_grad=True)  # weights with shape: (3, 2)
b = torch.ones(1, 2, requires_grad=True)  # bias with shape: (1, 2)
y = torch.ones(1, 2)                      # output with shape: (1, 2)

# forward pass
z = (x @ W) + b                           # pre-activation with shape: (1, 2)
y_pred = F.relu(z)                        # activation with shape: (1, 2)
loss = F.mse_loss(y_pred, y)              # scalar loss

叶子张量与非叶子张量#

运行前向传播后，PyTorch autograd 构建了一个动态计算图，如下图所示。这是一个有向无环图 (DAG)，它记录了输入张量（叶子节点）、对这些张量的所有后续操作以及中间/输出张量（非叶子节点）。计算图通过使用微积分中的链式法则，从图的根节点（输出）到叶子节点（输入）计算每个张量的梯度。

\[\mathbf{y} = \mathbf{f}_k\bigl(\mathbf{f}_{k-1}(\dots \mathbf{f}_1(\mathbf{x}) \dots)\bigr)\]

\[\frac{\partial \mathbf{y}}{\partial \mathbf{x}} = \frac{\partial \mathbf{f}_k}{\partial \mathbf{f}_{k-1}} \cdot \frac{\partial \mathbf{f}_{k-1}}{\partial \mathbf{f}_{k-2}} \cdot \cdots \cdot \frac{\partial \mathbf{f}_1}{\partial \mathbf{x}}\]

        graph TD

    x["x<br/>is_leaf=True<br/>requires_grad=False<br/>retains_grad=False<br/>grad=None"]
    W["W<br/>is_leaf=True<br/>requires_grad=True<br/>retains_grad=False<br/>grad=None"]
    b["b<br/>is_leaf=True<br/>requires_grad=True<br/>retains_grad=False<br/>grad=None"]
    matmul["x @ W"]
    z["z = x @ W + b<br/>is_leaf=False<br/>requires_grad=True<br/>retains_grad=False<br/>grad=None"]
    relu["y_pred = relu(z)<br/>is_leaf=False<br/>requires_grad=True<br/>retains_grad=False<br/>grad=None"]
    y["y<br/>is_leaf=True<br/>requires_grad=False<br/>retains_grad=False<br/>grad=None"]
    loss["loss = mse(y_pred, y)<br/>is_leaf=False<br/>requires_grad=True<br/>retains_grad=False<br/>grad=None"]

    x --> matmul
    W --> matmul
    matmul --> z
    b --> z
    z --> relu
    relu --> loss
    y --> loss

如果一个节点不是通过至少一个设置了 requires_grad=True 的输入进行张量运算的结果（例如 x, W, b 和 y），PyTorch 就会将其视为叶子 (leaf)；其他所有节点都被视为非叶子 (non-leaf)（例如 z, y_pred 和 loss）。您可以通过探测张量的 is_leaf 属性来以编程方式验证这一点。

# prints True because new tensors are leafs by convention
print(f"{x.is_leaf=}")

# prints False because tensor is the result of an operation with at
# least one input having requires_grad=True
print(f"{z.is_leaf=}")

x.is_leaf=True
z.is_leaf=False

叶子节点和非叶子节点的区别决定了在反向传播后，该张量的梯度是否会被存储在 grad 属性中，从而能否用于梯度下降。我们将在下一节中对此进行更多介绍。

现在让我们研究 PyTorch 如何计算并存储计算图中张量的梯度。

`requires_grad`#

为了构建可用于梯度计算的计算图，我们需要在张量构造函数中传入 requires_grad=True 参数。默认情况下，该值为 False，因此 PyTorch 不会跟踪任何已创建张量的梯度。要验证这一点，尝试不设置 requires_grad，重新运行前向传播，然后执行反向传播。您将看到

>>> loss.backward()
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

此错误意味着 autograd 无法反向传播到任何叶子张量，因为 loss 没有在跟踪梯度。如果您需要更改该属性，可以在张量上调用 requires_grad_()（注意末尾的下划线）。

我们可以像上面使用 is_leaf 属性那样，对哪些节点需要进行梯度计算进行完整性检查。

print(f"{x.requires_grad=}") # prints False because requires_grad=False by default
print(f"{W.requires_grad=}") # prints True because we set requires_grad=True in constructor
print(f"{z.requires_grad=}") # prints True because tensor is a non-leaf node

x.requires_grad=False
W.requires_grad=True
z.requires_grad=True

记住一点很有用：非叶子张量默认情况下 requires_grad=True，否则反向传播将会失败。如果张量是叶子张量，那么只有在用户明确设置的情况下，它才会拥有 requires_grad=True。另一种表述方式是：如果张量的至少一个输入需要梯度，那么该张量也会需要梯度。

该规则有两个例外：

任何具有 nn.Parameter 的 nn.Module，其参数的 requires_grad 都为 True（参见此处）。
使用上下文管理器局部禁用梯度计算（参见此处）。

总之，requires_grad 告诉 autograd 哪些张量在反向传播工作时需要计算其梯度。这与哪些张量的 grad 字段被填充是两回事，后者是下一节的主题。

`retain_grad`#

为了实际进行优化（例如 SGD、Adam 等），我们需要运行反向传播以便提取梯度。

loss.backward()

调用 backward() 会填充所有设置了 requires_grad=True 的叶子张量的 grad 字段。grad 是损失相对于我们要探测的张量的梯度。在运行 backward() 之前，此属性被设置为 None。

print(f"{W.grad=}")
print(f"{b.grad=}")

W.grad=tensor([[3., 3.],
        [3., 3.],
        [3., 3.]])
b.grad=tensor([[3., 3.]])

您可能对我们网络中的其他张量感到好奇。让我们检查剩余的叶子节点。

# prints all None because requires_grad=False
print(f"{x.grad=}")
print(f"{y.grad=}")

x.grad=None
y.grad=None

这些张量的梯度没有被填充，因为我们没有明确告知 PyTorch 计算它们的梯度（requires_grad=False）。

现在让我们看看一个中间的非叶子节点。

print(f"{z.grad=}")

/var/lib/workspace/beginner_source/understanding_leaf_vs_nonleaf_tutorial.py:230: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more information. (Triggered internally at /__w/pytorch/pytorch/build/aten/src/ATen/core/TensorBody.h:493.)
  print(f"{z.grad=}")
z.grad=None

PyTorch 为该梯度返回 None，并警告我们正在访问非叶子节点的 grad 属性。尽管 autograd 为了反向传播必须计算中间梯度，但它假定您事后不需要访问这些值。要改变这种行为，我们可以在张量上使用 retain_grad() 函数。这会告知 autograd 引擎在调用 backward() 后填充该张量的 grad。

# we have to re-run the forward pass
z = (x @ W) + b
y_pred = F.relu(z)
loss = F.mse_loss(y_pred, y)

# tell PyTorch to store the gradients after backward()
z.retain_grad()
y_pred.retain_grad()
loss.retain_grad()

# have to zero out gradients otherwise they would accumulate
W.grad = None
b.grad = None

# backpropagation
loss.backward()

# print gradients for all tensors that have requires_grad=True
print(f"{W.grad=}")
print(f"{b.grad=}")
print(f"{z.grad=}")
print(f"{y_pred.grad=}")
print(f"{loss.grad=}")

W.grad=tensor([[3., 3.],
        [3., 3.],
        [3., 3.]])
b.grad=tensor([[3., 3.]])
z.grad=tensor([[3., 3.]])
y_pred.grad=tensor([[3., 3.]])
loss.grad=tensor(1.)

我们得到与之前相同的 W.grad 结果。还要注意，因为 loss 是标量，损失相对于自身的梯度简单地为 1.0。

如果我们现在查看计算图的状态，会发现中间张量的 retains_grad 属性已经改变。按照惯例，对于任何叶子节点，此属性都会打印 False，即使它需要其梯度。

        graph TD

     x["x<br/>is_leaf=True<br/>requires_grad=False<br/>retains_grad=False<br/>grad=None"]
     W["W<br/>is_leaf=True<br/>requires_grad=True<br/>retains_grad=False<br/>grad=torch.Tensor"]
     b["b<br/>is_leaf=True<br/>requires_grad=True<br/>retains_grad=False<br/>grad=torch.Tensor"]
     matmul["x @ W"]
     z["z = x @ W + b<br/>is_leaf=False<br/>requires_grad=True<br/>retains_grad=True<br/>grad=torch.Tensor"]
     relu["y_pred = relu(z)<br/>is_leaf=False<br/>requires_grad=True<br/>retains_grad=True<br/>grad=torch.Tensor"]
     y["y<br/>is_leaf=True<br/>requires_grad=True<br/>retains_grad=False<br/>grad=None"]
     loss["loss = mse(y_pred, y)<br/>is_leaf=False<br/>requires_grad=True<br/>retains_grad=True<br/>grad=torch.Tensor"]

     x --> matmul
     W --> matmul
     matmul --> z
     b --> z
     z --> relu
     relu --> loss
     y --> loss

如果您在叶子张量上调用 retain_grad()，它不会执行任何操作，因为叶子张量默认情况下已经保留了它们的梯度（当 requires_grad=True 时）。如果我们对设置了 requires_grad=False 的张量调用 retain_grad()，PyTorch 实际上会抛出错误，因为它无法存储从未计算过的梯度。

>>> x.retain_grad()
RuntimeError: can't retain_grad on Tensor that has requires_grad=False

总结表#

使用 retain_grad() 和 retains_grad 仅对非叶子节点有意义，因为对于设置了 requires_grad=True 的叶子张量，其 grad 属性本身就已经会被填充。默认情况下，这些非叶子节点在反向传播后不会保留（存储）它们的梯度。我们可以通过重新运行前向传播、告知 PyTorch 存储梯度，然后执行反向传播来改变这一点。

下表可作为上述讨论的参考。以下场景是仅有的对 PyTorch 张量有效的情况。

`is_leaf (是否叶子节点)`	`requires_grad (是否需要梯度)`	`retains_grad (是否保留梯度)`	`require_grad()`	`retain_grad()`
`真`	`假`	`假`	将 `requires_grad` 设置为 `True` 或 `False`	抛出错误
`真`	`真`	`假`	将 `requires_grad` 设置为 `True` 或 `False`	无操作（已保留）
`假`	`真`	`假`	无操作	将 `retains_grad` 设置为 `True`
`假`	`真`	`真`	无操作	无操作（已保留）

结论#

在本教程中，我们介绍了 PyTorch 何时以及如何计算叶子和非叶子张量的梯度。通过使用 retain_grad，我们可以访问 autograd 计算图中中间张量的梯度。

如果您想了解更多关于 PyTorch autograd 系统的工作原理，请访问下方的参考文献。如果您对本教程有任何反馈（改进、更正错别字等），请使用 PyTorch 论坛和/或问题追踪器与我们联系。

参考文献#

脚本总运行时间: (0 分 0.366 秒)

理解 requires_grad, retain_grad, Leaf 和 Non-leaf 张量#

设置#

叶子张量与非叶子张量#

`requires_grad`#

`retain_grad`#

总结表#

结论#

参考文献#

文档

教程

资源

理解 requires_grad, retain_grad, Leaf 和 Non-leaf 张量#

设置#

叶子张量与非叶子张量#

requires_grad#

retain_grad#

总结表#

结论#

参考文献#

文档

教程

资源

`requires_grad`#

`retain_grad`#