PyTorch Autograd 引擎概述 – PyTorch

这篇博客文章基于 PyTorch 1.8 版本，尽管它也适用于旧版本，因为大部分机制保持不变。

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

什么是 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 符号的函数，以展示链式法则如何应用。


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

  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}$ 的梯度只有在乘法和正弦梯度相加后才完成。如您所见，我们计算了与 Jvp 等效的值，但没有构造矩阵。

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

PyTorch 自动微分引擎概述

什么是 autograd？

参考文献

文档

教程

资源

保持联系以获取更新、活动信息和最新新闻