PyTorch Autograd 引擎概述 – PyTorch

本篇博文基于 PyTorch 1.8 版本，尽管大多数机制保持不变，但内容也应适用于旧版本。

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

什么是自动求导（autograd）？

背景

PyTorch 通过使用自动微分（Automatic Differentiation）来计算函数相对于输入的梯度。自动微分是一种技术，它在给定计算图的情况下，计算输入的梯度。自动微分可以通过两种不同的方式执行：前向模式和反向模式。前向模式意味着我们在计算函数结果的同时计算梯度，而反向模式要求我们先对函数求值，然后从输出开始计算梯度。虽然这两种模式各有利弊，但由于输出数量通常少于输入数量，反向模式成为了事实上的首选，因为它允许进行更高效的计算。查看 [3] 以了解更多信息。

自动微分依赖于一个经典的微积分公式，即链式法则（chain-rule）。链式法则允许我们将复杂的导数拆分并随后重新组合，从而进行计算。

从形式上讲，给定复合函数 $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 中如何存储导数的一个例子。

假设我们当前正在处理 $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$ 反向步骤计算出的梯度。此步骤将一直重复，直到我们达到相对于输入的梯度，计算即告完成。 $\frac{\partial z}{x2}$ 的梯度仅在乘法和正弦梯度相加后才完成。如您所见，我们计算了等效于 Jvp 的结果，但没有构造整个矩阵。

在下一篇文章中，我们将深入研究 PyTorch 代码，看看这个计算图是如何构建的，以及如果您想尝试的话，相关部分在哪里！

PyTorch 自动微分引擎概述

什么是自动求导（autograd）？

参考文献

文档

教程

资源

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