梯度和自动微分简介

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

自动微分和梯度

自动微分 用于实现机器学习算法，例如用于训练神经网络的 反向传播。

在本指南中，您将探索使用 TensorFlow 计算梯度的方法，尤其是在急切执行中。

设置

import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf

计算梯度

为了自动微分，TensorFlow 需要记住在正向传递期间哪些操作按什么顺序执行。然后，在反向传递期间，TensorFlow 以相反的顺序遍历此操作列表以计算梯度。

梯度带

TensorFlow 提供了 tf.GradientTape API 用于自动微分；也就是说，计算某个目标（通常是损失）相对于某些源（通常是模型的变量）的梯度。TensorFlow 将 tf.GradientTape 上下文内执行的相关操作“记录”到“带”上。然后，TensorFlow 使用该带使用 反向模式微分 计算“记录”计算的梯度。

这是一个简单的示例

x = tf.Variable(3.0)

with tf.GradientTape() as tape:
  y = x**2

记录一些操作后，使用 GradientTape.gradient(target, sources) 计算某个目标（通常是损失）相对于某些源（通常是模型的变量）的梯度

# dy = 2x * dx
dy_dx = tape.gradient(y, x)
dy_dx.numpy()

上面的示例使用标量，但 tf.GradientTape 可以轻松地应用于任何张量

w = tf.Variable(tf.random.normal((3, 2)), name='w')
b = tf.Variable(tf.zeros(2, dtype=tf.float32), name='b')
x = [[1., 2., 3.]]

with tf.GradientTape(persistent=True) as tape:
  y = x @ w + b
  loss = tf.reduce_mean(y**2)

要获取 loss 相对于两个变量的梯度，可以将它们都作为源传递给 gradient 方法。磁带对源的传递方式很灵活，可以接受列表或字典的任何嵌套组合，并以相同的方式返回梯度（参见 tf.nest）。

[dl_dw, dl_db] = tape.gradient(loss, [w, b])

每个源的梯度具有与源相同的形状

print(w.shape)
print(dl_dw.shape)

以下是梯度计算，这次传递了一个变量字典

my_vars = {
    'w': w,
    'b': b
}

grad = tape.gradient(loss, my_vars)
grad['b']

模型的梯度

通常将 tf.Variables 收集到 tf.Module 或其子类（layers.Layer、keras.Model）中，以便进行 检查点 和 导出。

在大多数情况下，您需要计算模型可训练变量的梯度。由于 tf.Module 的所有子类都在 Module.trainable_variables 属性中聚合了它们的变量，因此您可以用几行代码计算这些梯度

layer = tf.keras.layers.Dense(2, activation='relu')
x = tf.constant([[1., 2., 3.]])

with tf.GradientTape() as tape:
  # Forward pass
  y = layer(x)
  loss = tf.reduce_mean(y**2)

# Calculate gradients with respect to every trainable variable
grad = tape.gradient(loss, layer.trainable_variables)

for var, g in zip(layer.trainable_variables, grad):
  print(f'{var.name}, shape: {g.shape}')

控制磁带观察的内容

默认行为是在访问可训练的 tf.Variable 后记录所有操作。这样做的原因是

• 磁带需要知道在正向传递中记录哪些操作，以便在反向传递中计算梯度。
• 磁带保存对中间输出的引用，因此您不希望记录不必要的操作。
• 最常见的用例是计算损失相对于模型所有可训练变量的梯度。

例如，以下代码无法计算梯度，因为 tf.Tensor 默认情况下不会被“观察”，并且 tf.Variable 不可训练

# A trainable variable
x0 = tf.Variable(3.0, name='x0')
# Not trainable
x1 = tf.Variable(3.0, name='x1', trainable=False)
# Not a Variable: A variable + tensor returns a tensor.
x2 = tf.Variable(2.0, name='x2') + 1.0
# Not a variable
x3 = tf.constant(3.0, name='x3')

with tf.GradientTape() as tape:
  y = (x0**2) + (x1**2) + (x2**2)

grad = tape.gradient(y, [x0, x1, x2, x3])

for g in grad:
  print(g)

可以使用 GradientTape.watched_variables 方法列出磁带正在观察的变量

[var.name for var in tape.watched_variables()]

tf.GradientTape 提供钩子，让用户可以控制观察哪些内容或不观察哪些内容。

要记录相对于 tf.Tensor 的梯度，需要调用 GradientTape.watch(x)

x = tf.constant(3.0)
with tf.GradientTape() as tape:
  tape.watch(x)
  y = x**2

# dy = 2x * dx
dy_dx = tape.gradient(y, x)
print(dy_dx.numpy())

相反，要禁用观察所有 tf.Variables 的默认行为，在创建梯度磁带时设置 watch_accessed_variables=False。此计算使用两个变量，但只连接其中一个变量的梯度

x0 = tf.Variable(0.0)
x1 = tf.Variable(10.0)

with tf.GradientTape(watch_accessed_variables=False) as tape:
  tape.watch(x1)
  y0 = tf.math.sin(x0)
  y1 = tf.nn.softplus(x1)
  y = y0 + y1
  ys = tf.reduce_sum(y)

由于没有对 x0 调用 GradientTape.watch，因此不会计算相对于它的梯度

# dys/dx1 = exp(x1) / (1 + exp(x1)) = sigmoid(x1)
grad = tape.gradient(ys, {'x0': x0, 'x1': x1})

print('dy/dx0:', grad['x0'])
print('dy/dx1:', grad['x1'].numpy())

中间结果

您还可以请求输出相对于 tf.GradientTape 上下文中计算的中间值的梯度。

x = tf.constant(3.0)

with tf.GradientTape() as tape:
  tape.watch(x)
  y = x * x
  z = y * y

# Use the tape to compute the gradient of z with respect to the
# intermediate value y.
# dz_dy = 2 * y and y = x ** 2 = 9
print(tape.gradient(z, y).numpy())

默认情况下，GradientTape 持有的资源会在调用 GradientTape.gradient 方法后立即释放。要对同一计算进行多次梯度计算，请使用 persistent=True 创建梯度磁带。这允许对 gradient 方法进行多次调用，因为资源会在磁带对象被垃圾回收时释放。例如

x = tf.constant([1, 3.0])
with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  y = x * x
  z = y * y

print(tape.gradient(z, x).numpy())  # [4.0, 108.0] (4 * x**3 at x = [1.0, 3.0])
print(tape.gradient(y, x).numpy())  # [2.0, 6.0] (2 * x at x = [1.0, 3.0])

del tape   # Drop the reference to the tape

性能注意事项

• 在梯度磁带上下文中执行操作会产生很小的开销。对于大多数急切执行来说，这不会是明显的成本，但您仍然应该只在需要的地方使用磁带上下文。

• 梯度磁带使用内存来存储中间结果，包括输入和输出，以便在反向传递期间使用。

  为了提高效率，某些操作（如 ReLU）不需要保留其中间结果，它们会在正向传递期间被修剪。但是，如果您在磁带上使用 persistent=True，则不会丢弃任何内容，并且您的峰值内存使用量会更高。

非标量目标的梯度

梯度本质上是对标量的操作。

x = tf.Variable(2.0)
with tf.GradientTape(persistent=True) as tape:
  y0 = x**2
  y1 = 1 / x

print(tape.gradient(y0, x).numpy())
print(tape.gradient(y1, x).numpy())

因此，如果您请求多个目标的梯度，则每个源的结果是

• 目标之和的梯度，或等效地
• 每个目标梯度的总和。

x = tf.Variable(2.0)
with tf.GradientTape() as tape:
  y0 = x**2
  y1 = 1 / x

print(tape.gradient({'y0': y0, 'y1': y1}, x).numpy())

类似地，如果目标不是标量，则会计算总和的梯度

x = tf.Variable(2.)

with tf.GradientTape() as tape:
  y = x * [3., 4.]

print(tape.gradient(y, x).numpy())

这使得对损失集合的总和或逐元素损失计算的总和求梯度变得很简单。

如果您需要每个项目的单独梯度，请参考 雅可比矩阵。

在某些情况下，您可以跳过雅可比矩阵。对于逐元素计算，总和的梯度给出了每个元素相对于其输入元素的导数，因为每个元素都是独立的

x = tf.linspace(-10.0, 10.0, 200+1)

with tf.GradientTape() as tape:
  tape.watch(x)
  y = tf.nn.sigmoid(x)

dy_dx = tape.gradient(y, x)

plt.plot(x, y, label='y')
plt.plot(x, dy_dx, label='dy/dx')
plt.legend()
_ = plt.xlabel('x')

控制流

因为梯度磁带会记录操作的执行顺序，所以 Python 控制流会自然地得到处理（例如，if 和 while 语句）。

这里在 if 的每个分支上使用不同的变量。梯度只连接到使用的变量

x = tf.constant(1.0)

v0 = tf.Variable(2.0)
v1 = tf.Variable(2.0)

with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  if x > 0.0:
    result = v0
  else:
    result = v1**2 

dv0, dv1 = tape.gradient(result, [v0, v1])

print(dv0)
print(dv1)

请记住，控制语句本身不可微分，因此它们对基于梯度的优化器不可见。

根据上述示例中 x 的值，磁带会记录 result = v0 或 result = v1**2。相对于 x 的梯度始终为 None。

dx = tape.gradient(result, x)

print(dx)

gradient 返回 None 的情况

当目标未连接到源时，gradient 将返回 None。

x = tf.Variable(2.)
y = tf.Variable(3.)

with tf.GradientTape() as tape:
  z = y * y
print(tape.gradient(z, x))

这里 z 显然没有连接到 x，但是有几种不那么明显的方式会导致梯度断开连接。

1. 用张量替换了变量

在关于 "控制磁带观察的内容" 的部分中，您看到磁带会自动观察 tf.Variable，但不会观察 tf.Tensor。

一个常见的错误是无意中用 tf.Tensor 替换了 tf.Variable，而不是使用 Variable.assign 来更新 tf.Variable。以下是一个示例

x = tf.Variable(2.0)

for epoch in range(2):
  with tf.GradientTape() as tape:
    y = x+1

  print(type(x).__name__, ":", tape.gradient(y, x))
  x = x + 1   # This should be `x.assign_add(1)`

2. 在 TensorFlow 之外进行了计算

如果计算退出 TensorFlow，磁带无法记录梯度路径。例如

x = tf.Variable([[1.0, 2.0],
                 [3.0, 4.0]], dtype=tf.float32)

with tf.GradientTape() as tape:
  x2 = x**2

  # This step is calculated with NumPy
  y = np.mean(x2, axis=0)

  # Like most ops, reduce_mean will cast the NumPy array to a constant tensor
  # using `tf.convert_to_tensor`.
  y = tf.reduce_mean(y, axis=0)

print(tape.gradient(y, x))

3. 对整数或字符串求梯度

整数和字符串不可微分。如果计算路径使用这些数据类型，则不会有梯度。

没有人期望字符串是可微分的，但是如果您没有指定 dtype，则很容易意外地创建 int 常量或变量。

x = tf.constant(10)

with tf.GradientTape() as g:
  g.watch(x)
  y = x * x

print(g.gradient(y, x))

TensorFlow 不会自动在类型之间进行转换，因此，在实践中，您通常会遇到类型错误而不是缺少梯度。

4. 对有状态对象求梯度

状态会阻止梯度。当您从有状态对象中读取时，磁带只能观察当前状态，而不能观察导致该状态的历史记录。

tf.Tensor 是不可变的。创建张量后，您无法更改它。它有值，但没有状态。到目前为止讨论的所有操作也是无状态的：tf.matmul 的输出只取决于其输入。

tf.Variable 具有内部状态，即其值。当您使用变量时，会读取状态。通常需要计算相对于变量的梯度，但变量的状态会阻止梯度计算进一步回溯。例如

x0 = tf.Variable(3.0)
x1 = tf.Variable(0.0)

with tf.GradientTape() as tape:
  # Update x1 = x1 + x0.
  x1.assign_add(x0)
  # The tape starts recording from x1.
  y = x1**2   # y = (x1 + x0)**2

# This doesn't work.
print(tape.gradient(y, x0))   #dy/dx0 = 2*(x1 + x0)

类似地，tf.data.Dataset 迭代器和 tf.queue 是有状态的，并且会阻止通过它们的张量上的所有梯度。

未注册梯度

某些 tf.Operation 已注册为不可微分，并且会返回 None。其他操作没有注册梯度。

tf.raw_ops 页面显示哪些低级操作已注册梯度。

如果您尝试对没有注册梯度的浮点操作求梯度，磁带将抛出错误，而不是静默地返回 None。这样您就知道出了问题。

例如，tf.image.adjust_contrast 函数包装了 raw_ops.AdjustContrastv2，它可能具有梯度，但梯度未实现

image = tf.Variable([[[0.5, 0.0, 0.0]]])
delta = tf.Variable(0.1)

with tf.GradientTape() as tape:
  new_image = tf.image.adjust_contrast(image, delta)

try:
  print(tape.gradient(new_image, [image, delta]))
  assert False   # This should not happen.
except LookupError as e:
  print(f'{type(e).__name__}: {e}')

如果您需要对该操作求微分，则需要实现梯度并注册它（使用 tf.RegisterGradient），或者使用其他操作重新实现该函数。

零而不是 None

在某些情况下，对于未连接的梯度，获得 0 而不是 None 会很方便。您可以使用 unconnected_gradients 参数决定在遇到未连接的梯度时返回什么

x = tf.Variable([2., 2.])
y = tf.Variable(3.)

with tf.GradientTape() as tape:
  z = y**2
print(tape.gradient(z, x, unconnected_gradients=tf.UnconnectedGradients.ZERO))