PyTorch Autograd 引擎概览

作者：Preferred Networks, Inc.

这篇博文基于 PyTorch 1.8 版本，尽管它也适用于较旧的版本，因为大多数机制保持不变。

为了帮助理解此处解释的概念，如果您不熟悉诸如 ATen 或 c10d 等 PyTorch 架构组件，建议您阅读 @ezyang 的精彩博文：PyTorch internals。

什么是 autograd？

背景

PyTorch 使用自动微分计算函数相对于输入的梯度。自动微分是一种技术，给定计算图，它计算输入的梯度。自动微分可以通过两种不同的方式执行：前向模式和反向模式。前向模式意味着我们在计算函数结果的同时计算梯度，而反向模式要求我们先评估函数，然后从输出开始计算梯度。虽然两种模式各有优缺点，但反向模式是事实上的选择，因为输出的数量小于输入的数量，这使得计算效率更高。请查阅 [3] 以了解更多信息。

自动微分依赖于一个经典的微积分公式，称为链式法则。链式法则允许我们将非常复杂的导数分解，然后再重新组合。

形式上，给定复合函数 $f(g(x))$ ，我们可以将其导数计算为 $\frac{\partial}{\partial x} f(g(x)) = f'(g(x))g'(x)$ 。这个结果使自动微分得以工作。通过组合构成更大函数（如神经网络）的更简单函数的导数，可以在给定点计算梯度的精确值，而不是依赖数值近似，后者需要多次扰动输入才能获得一个值。

为了直观理解反向模式的工作原理，让我们看一个简单的函数 $f(x, y) = log(x*y)$ 。图 1 显示了其计算图，其中左侧的输入 x、y 通过一系列操作流向输出 z。

图 1：f(x, y) = log(x*y) 的计算图

自动微分引擎通常会执行此图。它还会扩展此图以计算 w 相对于输入 x、y 以及中间结果 v 的导数。

上述示例函数可以分解为 f 和 g，其中 $f(x, y) = log(g(x, y))$ 且 $g(x, y) = xy$ 。每当引擎执行图中的一个操作时，该操作的导数就会添加到图中以便稍后在反向传播中执行。注意，引擎知道基本函数的导数。

在上面的示例中，当将 x 和 y 相乘得到 v 时，引擎将使用它已知的乘法导数定义扩展图，以计算乘法的偏导数。 $\frac{\partial}{\partial x} g(x, y) = y$ 和 $\frac{\partial}{\partial y} g(x, y) = x$ 。结果扩展后的图如图 2 所示，其中 MultDerivative 节点还通过输入梯度计算结果梯度的乘积以应用链式法则；这将在后续操作中明确看到。请注意，反向图（绿色节点）直到所有前向步骤完成后才会执行。

图 2：执行对数操作后扩展的计算图

接着，引擎现在计算 $log(v)$ 操作，并再次使用它已知的对数导数（即 $\frac{1}{v}$ ）扩展图。这显示在图 3 中。此操作生成结果 $\frac{\partial w}{\partial v}$ ，当其反向传播并按照链式法则乘以乘法导数时，生成导数 $\frac{\partial w}{\partial x}$ 和 $\frac{\partial w}{\partial y}$ 。

图 3：执行对数操作后扩展的计算图

原始计算图被扩展了一个新的哑变量 z，它与 w 相同。z 相对于 w 的导数是 1，因为它们是同一个变量，这个技巧允许我们应用链式法则来计算输入的导数。在前向传播完成后，我们开始反向传播，为 $\frac{\partial z}{\partial w}$ 提供初始值 1.0。这显示在图 4 中。

图 4：为反向自动微分扩展的计算图

然后沿着绿色图执行自动微分引擎引入的 LogDerivative 操作 $\frac{1}{v}$ ，并将其结果乘以 $\frac{\partial z}{\partial w}$ ，以根据链式法则获得梯度 $\frac{\partial z}{\partial v}$ 。接下来，以相同的方式执行乘法导数，最终获得所需的导数 $\frac{\partial z}{\partial x} \frac{\partial z}{\partial y}$ 。

形式上，我们在这里以及 PyTorch autograd 引擎所做的是计算雅可比向量积 (Jvp)，以计算模型参数的梯度，因为模型参数和输入都是向量。

雅可比向量积

当我们计算向量值函数 $f(\overline{x}) = \overline{y}$ （输入和输出都是向量的函数）的梯度时，我们实质上是构建一个雅可比矩阵。

多亏了链式法则，将函数 $f(\overline{x}) = \overline{y}$ 的雅可比矩阵乘以向量 $v$ ，其中 $v$ 是先前计算的标量函数 $z = g(\overline{y})$ 的梯度，结果是标量输出相对于向量值函数输入的梯度 $\frac{\partial z}{\partial x_1} \cdots \frac{\partial z}{\partial x_n}$ 。

例如，让我们看看一些用 python 符号表示的函数，以展示链式法则如何应用。

$f(x_1, x_2) = (log(x_1 x_2), sin(x_2))$

def f(x1, x2):
      a = x1 * x2
      y1 = log(a)
      y2 = sin(x2)
      return (y1, y2)

$g(y_1, y_2) = y_1 y_2$

def g(y1, y2):
      return y1 * y2

现在，如果我们使用链式法则和导数定义手工推导，我们将获得以下一组恒等式，我们可以直接将其代入 $f(x_1, x_2)$ 的雅可比矩阵

$\frac{\partial y_1}{\partial x_1} = \frac{\partial y_1}{\partial a}\frac{\partial a}{\partial x_1} = \frac{1}{x_1}$

$\frac{\partial y_1}{\partial x_2} = \frac{\partial y_1}{\partial a}\frac{\partial a}{\partial x_2} = \frac{1}{x_2}$

$\frac{\partial y_2}{\partial x_1} = 0$

$\frac{\partial y_2}{\partial x_2} = cos(x_2)$

接下来，考虑标量函数 $z = g(y_1, y_2)$ 的梯度

$\frac{\partial z}{y_1} = y_2$

$\frac{\partial z}{y_2} = y_1$

如果现在按照链式法则计算转置雅可比向量积，我们得到以下表达式

$\begin{pmatrix}\frac{\partial y_1}{\partial x_1} & \frac{\partial y_1}{\partial x_2} \\\\ \frac{\partial y_2}{\partial x_1} & \frac{\partial y_2}{\partial x_2} \end{pmatrix}^{t} \begin{pmatrix} y_2\\y_1 \end{pmatrix} = \begin{pmatrix}\frac{1}{x_1} & \frac{1}{x_2} \\\\ 0 & cos(x_2)) \end{pmatrix}^{t} \begin{pmatrix} y_2\\y_1 \end{pmatrix} = \begin{pmatrix} \frac{1}{x_1}y_2\\\frac{1}{x_2}y_2+cos(x_2)y_1 \end{pmatrix}$

计算 $(x_1, x_2) = (0.5, 0.75)$ 的 Jvp 得出结果： $(\frac{dy}{x_1}, \frac{dy}{x_2}) = (1.3633, 0.1912)$ 我们可以在 PyTorch 中执行相同的表达式并计算输入的梯度

>>> import torch

>>> x = torch.tensor([0.5, 0.75], requires_grad=True)

>>> y = torch.log(x[0] * x[1]) * torch.sin(x[1])

>>> y.backward(1.0)

>>> x.grad

tensor([1.3633, 0.1912])

结果与我们手工计算的雅可比向量积相同！然而，PyTorch 从未构建矩阵，因为它可能会变得非常大，而是创建了一个操作图，该图在反向遍历时应用了 tools/autograd/derivatives.yaml 中定义的雅可比向量积。

遍历计算图

每当 PyTorch 执行一个操作时，autograd 引擎都会构建一个用于反向遍历的图。反向模式自动微分首先在末尾添加一个标量变量 $z = w$ ，以便 $\frac{\partial z}{\partial w} = 1$ ，正如我们在引言中所见。这是提供给 Jvp 引擎计算的初始梯度值，正如我们在上一节中所见。

在 PyTorch 中，用户在调用 backward 方法时会显式设置初始梯度。

然后，Jvp 计算开始，但它从不构建矩阵。相反，当 PyTorch 记录计算图时，会添加已执行前向操作的导数（反向节点）。图 5 显示了执行前面看到的函数 $f(x_1, x_2)$ 和 $g(y_1, y_2)$ 所生成的反向图。

图 5：扩展了反向传播的计算图

前向传播完成后，结果用于反向传播，其中执行计算图中的导数。基本导数存储在 tools/autograd/derivatives.yaml 文件中，它们不是常规导数，而是它们的 Jvp 版本 [3]。它们将原始函数的输入和输出以及函数输出相对于最终输出的梯度作为参数。通过反复将结果梯度乘以图中下一个 Jvp 导数，将按照链式法则生成直至输入的梯度。

图 6：链式法则在反向微分中的应用方式

图 6 通过展示链式法则来表示此过程。我们以 1.0 作为初始值，正如前面所详细说明的，这是已经计算出的梯度 $\frac{\partial y}{\partial u}$ ，在图中以绿色高亮显示。然后我们移动到图中的下一个节点。在 derivatives.yaml 中注册的 backward 函数将计算相关的 $\frac{\partial u}{\partial v}$ 值（在图中以红色高亮显示），并将其乘以 $\frac{\partial y}{\partial u}$ 。根据链式法则，结果是 $\frac{\partial y}{\partial v}$ ，当我们在图中处理下一个反向节点时，这将成为已经计算出的梯度（绿色）。

您可能还注意到，在图 5 中，有一个梯度来自两个不同的源。当两个不同的函数共享一个输入时，该输入相对于输出的梯度会被累加，并且除非所有路径都已累加在一起，否则使用该梯度的计算无法继续。

让我们看一个 PyTorch 中导数存储方式的示例。

假设我们当前正在处理图 2 中 LogBackward 节点中的 $log$ 函数的反向传播。在 derivatives.yaml 中指定的 $log$ 的导数是 grad.div(self.conj())。grad 是已经计算出的梯度 $\frac{\partial z}{\partial y1}$ ，而 self.conj() 是输入向量的复共轭。对于复数，PyTorch 计算一个特殊的导数，称为共轭 Wirtinger 导数 [6]。这个导数接收复数及其共轭，并通过 [6] 中描述的一些神奇操作，它们在插入优化器时是最速下降的方向。

这段代码转换为 $(\frac{\partial z}{\partial y1} \frac{1}{v})$ ，对应于图 3 中的绿色和红色方块。继续，autograd 引擎将执行下一个操作：乘法的反向传播。与之前一样，输入是原始函数的输入以及从 $log$ 反向步骤计算出的梯度。此步骤将重复进行，直到我们达到相对于输入的梯度，并且计算将完成。 $\frac{\partial z}{x2}$ 的梯度只有在乘法和 sin 梯度相加后才完成。正如您所见，我们计算了 Jvp 的等价物，但没有构建矩阵。

在下一篇博文中，我们将深入 PyTorch 代码，看看这个图是如何构建的，以及如果您想对其进行实验，哪些是相关的部分！

参考文献

推荐：展示了反向传播如何在形式上用雅可比矩阵表示

https://cs.ubc.ca/~fwood/CS340/lectures/AD1.pdf

什么是 autograd？

参考文献

文档

教程

资源