2.5. 自动求导¶ 在 SageMaker Studio Lab 中打开 Notebook
回想一下 第 2.4 节，计算导数是我们用来训练深度网络的所有优化算法中的关键步骤。虽然计算本身很简单，但手动计算可能既繁琐又容易出错，而且随着模型变得越来越复杂，这些问题只会加剧。
幸运的是，所有现代深度学习框架都通过提供*自动求导*（通常简称为 *autograd*）为我们解决了这个问题。当我们将数据通过每个连续的函数传递时，框架会构建一个*计算图*，该图跟踪每个值如何依赖于其他值。为了计算导数，自动求导通过应用链式法则反向遍历此图。以这种方式应用链式法则的计算算法称为*反向传播*。
虽然 autograd 库在过去十年中已成为热门话题，但它们有着悠久的历史。事实上，最早提到 autograd 的文献可以追溯到半个多世纪前 (Wengert, 1964)。现代反向传播背后的核心思想可以追溯到 1980 年的一篇博士论文 (Speelpenning, 1980)，并在 1980 年代后期得到进一步发展 (Griewank, 1989)。虽然反向传播已成为计算梯度的默认方法，但它并不是唯一的选择。例如，Julia 编程语言采用正向传播 (Revels et al., 2016)。在探索各种方法之前，让我们先掌握 autograd 包。
import torch
from mxnet import autograd, np, npx
npx.set_np()
from jax import numpy as jnp
import tensorflow as tf
2.5.1. 一个简单的函数¶
假设我们有兴趣对函数 \(y = 2\mathbf{x}^{\top}\mathbf{x}\) 关于列向量 \(\mathbf{x}\) 求导。首先，我们给 x 赋一个初始值。
x = torch.arange(4.0)
x
tensor([0., 1., 2., 3.])
在我们计算 \(y\) 关于 \(\mathbf{x}\) 的梯度之前，我们需要一个地方来存储它。通常，我们避免在每次求导时都分配新的内存，因为深度学习需要连续多次地对相同参数计算导数，我们可能会有耗尽内存的风险。请注意，一个标量值函数关于一个向量 \(\mathbf{x}\) 的梯度是一个与 \(\mathbf{x}\) 形状相同的向量值。
# Can also create x = torch.arange(4.0, requires_grad=True)
x.requires_grad_(True)
x.grad # The gradient is None by default
x = np.arange(4.0)
x
[22:07:05] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for CPU
array([0., 1., 2., 3.])
在我们计算 \(y\) 关于 \(\mathbf{x}\) 的梯度之前，我们需要一个地方来存储它。通常，我们避免在每次求导时都分配新的内存，因为深度学习需要连续多次地对相同参数计算导数，我们可能会有耗尽内存的风险。请注意，一个标量值函数关于一个向量 \(\mathbf{x}\) 的梯度是一个与 \(\mathbf{x}\) 形状相同的向量值。
# We allocate memory for a tensor's gradient by invoking `attach_grad`
x.attach_grad()
# After we calculate a gradient taken with respect to `x`, we will be able to
# access it via the `grad` attribute, whose values are initialized with 0s
x.grad
array([0., 0., 0., 0.])
x = jnp.arange(4.0)
x
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Array([0., 1., 2., 3.], dtype=float32)
x = tf.range(4, dtype=tf.float32)
x
<tf.Tensor: shape=(4,), dtype=float32, numpy=array([0., 1., 2., 3.], dtype=float32)>
在我们计算 \(y\) 关于 \(\mathbf{x}\) 的梯度之前，我们需要一个地方来存储它。通常，我们避免在每次求导时都分配新的内存，因为深度学习需要连续多次地对相同参数计算导数，我们可能会有耗尽内存的风险。请注意，一个标量值函数关于一个向量 \(\mathbf{x}\) 的梯度是一个与 \(\mathbf{x}\) 形状相同的向量值。
x = tf.Variable(x)
现在我们计算我们的函数 x 并将结果赋给 y。
y = 2 * torch.dot(x, x)
y
tensor(28., grad_fn=<MulBackward0>)
我们现在可以通过调用其 backward 方法来求 y 关于 x 的梯度。接下来，我们可以通过 x 的 grad 属性访问梯度。
y.backward()
x.grad
tensor([ 0., 4., 8., 12.])
# Our code is inside an `autograd.record` scope to build the computational
# graph
with autograd.record():
y = 2 * np.dot(x, x)
y
array(28.)
我们现在可以通过调用其 backward 方法来求 y 关于 x 的梯度。接下来，我们可以通过 x 的 grad 属性访问梯度。
y.backward()
x.grad
[22:07:05] ../src/base.cc:48: GPU context requested, but no GPUs found.
array([ 0., 4., 8., 12.])
y = lambda x: 2 * jnp.dot(x, x)
y(x)
Array(28., dtype=float32)
我们现在可以通过传递 grad 变换来计算 y 关于 x 的梯度。
from jax import grad
# The `grad` transform returns a Python function that
# computes the gradient of the original function
x_grad = grad(y)(x)
x_grad
Array([ 0., 4., 8., 12.], dtype=float32)
# Record all computations onto a tape
with tf.GradientTape() as t:
y = 2 * tf.tensordot(x, x, axes=1)
y
<tf.Tensor: shape=(), dtype=float32, numpy=28.0>
我们现在可以通过调用 gradient 方法来计算 y 关于 x 的梯度。
x_grad = t.gradient(y, x)
x_grad
<tf.Tensor: shape=(4,), dtype=float32, numpy=array([ 0., 4., 8., 12.], dtype=float32)>
我们已经知道函数 \(y = 2\mathbf{x}^{\top}\mathbf{x}\) 关于 \(\mathbf{x}\) 的梯度应该是 \(4\mathbf{x}\)。我们现在可以验证自动梯度计算和预期结果是否相同。
x.grad == 4 * x
tensor([True, True, True, True])
现在让我们计算另一个关于 x 的函数并求其梯度。请注意，PyTorch 在我们记录新梯度时不会自动重置梯度缓冲区。相反，新的梯度会加到已存储的梯度上。当我们想要优化多个目标函数的和时，这种行为很有用。要重置梯度缓冲区，我们可以如下调用 x.grad.zero_()。
x.grad.zero_() # Reset the gradient
y = x.sum()
y.backward()
x.grad
tensor([1., 1., 1., 1.])
x.grad == 4 * x
array([ True, True, True, True])
现在让我们计算另一个关于 x 的函数并求其梯度。注意，MXNet 在我们每次记录新梯度时都会重置梯度缓冲区。
with autograd.record():
y = x.sum()
y.backward()
x.grad # Overwritten by the newly calculated gradient
array([1., 1., 1., 1.])
x_grad == 4 * x
Array([ True, True, True, True], dtype=bool)
y = lambda x: x.sum()
grad(y)(x)
Array([1., 1., 1., 1.], dtype=float32)
x_grad == 4 * x
<tf.Tensor: shape=(4,), dtype=bool, numpy=array([ True, True, True, True])>
现在让我们计算另一个关于 x 的函数并求其梯度。注意，TensorFlow 在我们每次记录新梯度时都会重置梯度缓冲区。
with tf.GradientTape() as t:
y = tf.reduce_sum(x)
t.gradient(y, x) # Overwritten by the newly calculated gradient
<tf.Tensor: shape=(4,), dtype=float32, numpy=array([1., 1., 1., 1.], dtype=float32)>
2.5.2. 非标量变量的反向传播¶
当 y 是一个向量时，y 关于向量 x 的导数最自然的表示是一个称为*雅可比矩阵*的矩阵，它包含了 y 的每个分量关于 x 的每个分量的偏导数。同样，对于更高阶的 y 和 x，求导的结果可能是一个更高阶的张量。
虽然雅可比矩阵确实在一些高级机器学习技术中出现，但更常见的情况是，我们希望将 y 的每个分量的梯度相对于整个向量 x 求和，从而得到一个与 x 形状相同的向量。例如，我们经常有一个向量，表示在一*批*训练样本中为每个样本单独计算的损失函数值。这里，我们只想对每个样本单独计算的梯度求和。
由于深度学习框架对非标量张量的梯度解释各不相同，PyTorch 采取了一些措施来避免混淆。在非标量上调用 backward 会引发错误，除非我们告诉 PyTorch 如何将该对象约简为标量。更正式地说，我们需要提供某个向量 \(\mathbf{v}\)，这样 backward 将计算 \(\mathbf{v}^\top \partial_{\mathbf{x}} \mathbf{y}\) 而不是 \(\partial_{\mathbf{x}} \mathbf{y}\)。这部分可能令人困惑，但出于后面会阐明的原因，这个参数（代表 \(\mathbf{v}\)）被命名为 gradient。有关更详细的描述，请参见 Yang Zhang 的 Medium 帖子。
x.grad.zero_()
y = x * x
y.backward(gradient=torch.ones(len(y))) # Faster: y.sum().backward()
x.grad
tensor([0., 2., 4., 6.])
MXNet 通过在计算梯度之前将所有张量求和来约简为标量来处理这个问题。换句话说，它不是返回雅可比矩阵 \(\partial_{\mathbf{x}} \mathbf{y}\)，而是返回和的梯度 \(\partial_{\mathbf{x}} \sum_i y_i\)。
with autograd.record():
y = x * x
y.backward()
x.grad # Equals the gradient of y = sum(x * x)
array([0., 2., 4., 6.])
y = lambda x: x * x
# grad is only defined for scalar output functions
grad(lambda x: y(x).sum())(x)
Array([0., 2., 4., 6.], dtype=float32)
默认情况下，TensorFlow 返回和的梯度。换句话说，它不是返回雅可比矩阵 \(\partial_{\mathbf{x}} \mathbf{y}\)，而是返回和的梯度 \(\partial_{\mathbf{x}} \sum_i y_i\)。
with tf.GradientTape() as t:
y = x * x
t.gradient(y, x) # Same as y = tf.reduce_sum(x * x)
<tf.Tensor: shape=(4,), dtype=float32, numpy=array([0., 2., 4., 6.], dtype=float32)>
2.5.3. 分离计算¶
有时，我们希望将某些计算移出记录的计算图。例如，假设我们使用输入来创建一些辅助中间项，而我们不希望为这些项计算梯度。在这种情况下，我们需要将相应的计算图与最终结果*分离*。下面的玩具示例更清楚地说明了这一点：假设我们有 z = x * y 和 y = x * x，但我们希望关注 x 对 z 的*直接*影响，而不是通过 y 传递的影响。在这种情况下，我们可以创建一个新变量 u，它与 y 的值相同，但其*出处*（它是如何创建的）已被清除。因此，u 在图中没有祖先，梯度不会通过 u 流向 x。例如，对 z = x * u 求梯度将得到结果 u，（而不是你可能期望的 3 * x * x，因为 z = x * x * x）。
x.grad.zero_()
y = x * x
u = y.detach()
z = u * x
z.sum().backward()
x.grad == u
tensor([True, True, True, True])
with autograd.record():
y = x * x
u = y.detach()
z = u * x
z.backward()
x.grad == u
array([ True, True, True, True])
import jax
y = lambda x: x * x
# jax.lax primitives are Python wrappers around XLA operations
u = jax.lax.stop_gradient(y(x))
z = lambda x: u * x
grad(lambda x: z(x).sum())(x) == y(x)
Array([ True, True, True, True], dtype=bool)
# Set persistent=True to preserve the compute graph.
# This lets us run t.gradient more than once
with tf.GradientTape(persistent=True) as t:
y = x * x
u = tf.stop_gradient(y)
z = u * x
x_grad = t.gradient(z, x)
x_grad == u
<tf.Tensor: shape=(4,), dtype=bool, numpy=array([ True, True, True, True])>
请注意，虽然此过程将 y 的祖先从通向 z 的图中分离出来，但通向 y 的计算图仍然存在，因此我们可以计算 y 关于 x 的梯度。
x.grad.zero_()
y.sum().backward()
x.grad == 2 * x
tensor([True, True, True, True])
y.backward()
x.grad == 2 * x
array([ True, True, True, True])
grad(lambda x: y(x).sum())(x) == 2 * x
Array([ True, True, True, True], dtype=bool)
t.gradient(y, x) == 2 * x
<tf.Tensor: shape=(4,), dtype=bool, numpy=array([ True, True, True, True])>
2.5.4. 梯度和 Python 控制流¶
到目前为止，我们回顾了从输入到输出的路径通过像 z = x * x * x 这样的函数明确定义的情况。编程在如何计算结果方面为我们提供了更多的自由。例如，我们可以让它们依赖于辅助变量，或根据中间结果进行条件选择。使用自动求导的一个好处是，即使构建一个函数的计算图需要通过一系列 Python 控制流（例如，条件、循环和任意函数调用），我们仍然可以计算结果变量的梯度。为了说明这一点，请考虑以下代码片段，其中 while 循环的迭代次数和 if 语句的求值都取决于输入 a 的值。
def f(a):
b = a * 2
while b.norm() < 1000:
b = b * 2
if b.sum() > 0:
c = b
else:
c = 100 * b
return c
def f(a):
b = a * 2
while np.linalg.norm(b) < 1000:
b = b * 2
if b.sum() > 0:
c = b
else:
c = 100 * b
return c
def f(a):
b = a * 2
while jnp.linalg.norm(b) < 1000:
b = b * 2
if b.sum() > 0:
c = b
else:
c = 100 * b
return c
def f(a):
b = a * 2
while tf.norm(b) < 1000:
b = b * 2
if tf.reduce_sum(b) > 0:
c = b
else:
c = 100 * b
return c
下面，我们调用这个函数，传入一个随机值作为输入。由于输入是一个随机变量，我们不知道计算图将采取何种形式。但是，每当我们在特定输入上执行 f(a) 时，我们都会实现一个特定的计算图，并随后可以运行 backward。
a = torch.randn(size=(), requires_grad=True)
d = f(a)
d.backward()
a = np.random.normal()
a.attach_grad()
with autograd.record():
d = f(a)
d.backward()
from jax import random
a = random.normal(random.PRNGKey(1), ())
d = f(a)
d_grad = grad(f)(a)
a = tf.Variable(tf.random.normal(shape=()))
with tf.GradientTape() as t:
d = f(a)
d_grad = t.gradient(d, a)
d_grad
<tf.Tensor: shape=(), dtype=float32, numpy=2048.0>
尽管我们的函数 f 为了演示目的有点刻意，但它对输入的依赖非常简单：它是 a 的一个*线性*函数，具有分段定义的尺度。因此，f(a) / a 是一个常数条目的向量，而且 f(a) / a 需要与 f(a) 关于 a 的梯度相匹配。
a.grad == d / a
tensor(True)
a.grad == d / a
array(True)
d_grad == d / a
Array(True, dtype=bool)
d_grad == d / a
<tf.Tensor: shape=(), dtype=bool, numpy=True>
动态控制流在深度学习中非常常见。例如，在处理文本时，计算图取决于输入的长度。在这些情况下，自动求导对于统计建模至关重要，因为*先验*地计算梯度是不可能的。
2.5.5. 讨论¶
您现在已经领略了自动求导的威力。自动且高效地计算导数的库的开发极大地提高了深度学习从业者的生产力，使他们能够专注于不那么琐碎的工作。此外，autograd 让我们能够设计出庞大的模型，对于这些模型，用笔和纸计算梯度将是极其耗时的。有趣的是，虽然我们使用 autograd 来（在统计意义上）*优化*模型，但 autograd 库本身的*优化*（在计算意义上）是框架设计者极为关注的一个丰富课题。在这里，编译器和图操作的工具被用来以最快速和最节省内存的方式计算结果。
现在，请尝试记住这些基础知识：（i）将梯度附加到我们希望求导的那些变量上；（ii）记录目标值的计算过程；（iii）执行反向传播函数；以及（iv）访问得到的梯度。
2.5.6. 练习¶
为什么二阶导数的计算比一阶导数昂贵得多？
运行反向传播函数后，立即再次运行它，看看会发生什么。调查一下。
在控制流示例中，我们计算
d关于a的导数，如果我们将变量a更改为随机向量或矩阵会发生什么？此时，计算f(a)的结果不再是标量。结果会发生什么变化？我们如何分析这个问题？设 \(f(x) = \sin(x)\)。绘制 \(f\) 及其导数 \(f'\) 的图像。不要利用 \(f'(x) = \cos(x)\) 这一事实，而是使用自动求导来得到结果。
设 \(f(x) = ((\log x^2) \cdot \sin x) + x^{-1}\)。写出从 \(x\) 到 \(f(x)\) 的依赖关系图。
使用链式法则计算上述函数的导数 \(\frac{df}{dx}\)，并将每个项放在您先前构建的依赖关系图上。
给定图和中间导数结果，您在计算梯度时有多种选择。一次从 \(x\) 到 \(f\) 评估结果，一次从 \(f\) 回溯到 \(x\)。从 \(x\) 到 \(f\) 的路径通常称为*正向求导*，而从 \(f\) 到 \(x\) 的路径称为反向求导。
什么时候你可能想使用正向求导，什么时候使用反向求导？提示：考虑所需的中间数据量、并行化步骤的能力以及所涉及的矩阵和向量的大小。