PyTorch Autograd 引擎概述

作者：Preferred Networks, Inc.

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

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

什么是 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$ 与先前计算的标量函数 $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 中存储的示例。

假设我们当前正在处理 $log$ 函数的反向传播，在图 2 中的 LogBackward 节点中。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$ 反向步骤计算的梯度。此步骤将不断重复，直到我们获得相对于输入的梯度，并且计算将完成。只有当乘法和 sin 梯度加在一起时， $\frac{\partial z}{x2}$ 的梯度才算完成。如您所见，我们计算了 Jvp 的等效值，但没有构造矩阵。

在下一篇文章中，我们将深入 PyTorch 代码，了解如何构造此图以及如果您想尝试它，相关的部分在哪里！

参考文献

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

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

什么是 autograd？

参考文献

文档

教程

资源