高级自动微分

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

该 梯度和自动微分简介 指南包含在 TensorFlow 中计算梯度所需的一切。本指南重点介绍 tf.GradientTape API 的更深入、不太常见的特性。

设置

import tensorflow as tf

import matplotlib as mpl
import matplotlib.pyplot as plt

mpl.rcParams['figure.figsize'] = (8, 6)

控制梯度记录

在 自动微分指南 中，您了解了如何在构建梯度计算时控制哪些变量和张量被磁带监视。

磁带还具有用于操作记录的方法。

停止记录

如果您希望停止记录梯度，可以使用 tf.GradientTape.stop_recording 暂时暂停记录。

这可能有助于减少开销，如果您不希望在模型中间对复杂操作进行微分。这可能包括计算指标或中间结果

x = tf.Variable(2.0)
y = tf.Variable(3.0)

with tf.GradientTape() as t:
  x_sq = x * x
  with t.stop_recording():
    y_sq = y * y
  z = x_sq + y_sq

grad = t.gradient(z, {'x': x, 'y': y})

print('dz/dx:', grad['x'])  # 2*x => 4
print('dz/dy:', grad['y'])

重置/从头开始记录

如果您希望完全重新开始，请使用 tf.GradientTape.reset。退出梯度磁带块并重新启动通常更容易阅读，但当退出磁带块很困难或不可能时，您可以使用 reset 方法。

x = tf.Variable(2.0)
y = tf.Variable(3.0)
reset = True

with tf.GradientTape() as t:
  y_sq = y * y
  if reset:
    # Throw out all the tape recorded so far.
    t.reset()
  z = x * x + y_sq

grad = t.gradient(z, {'x': x, 'y': y})

print('dz/dx:', grad['x'])  # 2*x => 4
print('dz/dy:', grad['y'])

精确地停止梯度流

与上面的全局磁带控制相反，tf.stop_gradient 函数更加精确。它可以用于停止梯度沿着特定路径流动，而无需访问磁带本身

x = tf.Variable(2.0)
y = tf.Variable(3.0)

with tf.GradientTape() as t:
  y_sq = y**2
  z = x**2 + tf.stop_gradient(y_sq)

grad = t.gradient(z, {'x': x, 'y': y})

print('dz/dx:', grad['x'])  # 2*x => 4
print('dz/dy:', grad['y'])

自定义梯度

在某些情况下，您可能希望精确控制梯度的计算方式，而不是使用默认方式。这些情况包括

1. 您正在编写的新的操作没有定义的梯度。
2. 默认计算在数值上不稳定。
3. 您希望缓存正向传递中的昂贵计算。
4. 您想要修改一个值（例如，使用 tf.clip_by_value 或 tf.math.round）而不修改梯度。

对于第一种情况，要编写一个新的操作，您可以使用 tf.RegisterGradient 来设置自己的操作（有关详细信息，请参阅 API 文档）。（请注意，梯度注册表是全局的，因此请谨慎更改。）

对于后三种情况，您可以使用 tf.custom_gradient。

以下是一个将 tf.clip_by_norm 应用于中间梯度的示例

# Establish an identity operation, but clip during the gradient pass.
@tf.custom_gradient
def clip_gradients(y):
  def backward(dy):
    return tf.clip_by_norm(dy, 0.5)
  return y, backward

v = tf.Variable(2.0)
with tf.GradientTape() as t:
  output = clip_gradients(v * v)
print(t.gradient(output, v))  # calls "backward", which clips 4 to 2

有关更多详细信息，请参阅 tf.custom_gradient 装饰器 API 文档。

SavedModel 中的自定义梯度

自定义梯度可以通过使用选项 tf.saved_model.SaveOptions(experimental_custom_gradients=True) 保存到 SavedModel 中。

要保存到 SavedModel 中，梯度函数必须是可追踪的（要了解更多信息，请查看 使用 tf.function 提高性能 指南）。

class MyModule(tf.Module):

  @tf.function(input_signature=[tf.TensorSpec(None)])
  def call_custom_grad(self, x):
    return clip_gradients(x)

model = MyModule()

tf.saved_model.save(
    model,
    'saved_model',
    options=tf.saved_model.SaveOptions(experimental_custom_gradients=True))

# The loaded gradients will be the same as the above example.
v = tf.Variable(2.0)
loaded = tf.saved_model.load('saved_model')
with tf.GradientTape() as t:
  output = loaded.call_custom_grad(v * v)
print(t.gradient(output, v))

关于以上示例的说明：如果您尝试将以上代码替换为 tf.saved_model.SaveOptions(experimental_custom_gradients=False)，梯度在加载时仍将产生相同的结果。原因是梯度注册表仍然包含函数 call_custom_op 中使用的自定义梯度。但是，如果您在保存没有自定义梯度的情况下重新启动运行时，则在 tf.GradientTape 下运行加载的模型将抛出错误：LookupError: No gradient defined for operation 'IdentityN' (op type: IdentityN)。

多个磁带

多个磁带可以无缝交互。

例如，这里每个磁带都观察一组不同的张量

x0 = tf.constant(0.0)
x1 = tf.constant(0.0)

with tf.GradientTape() as tape0, tf.GradientTape() as tape1:
  tape0.watch(x0)
  tape1.watch(x1)

  y0 = tf.math.sin(x0)
  y1 = tf.nn.sigmoid(x1)

  y = y0 + y1

  ys = tf.reduce_sum(y)

tape0.gradient(ys, x0).numpy()   # cos(x) => 1.0

tape1.gradient(ys, x1).numpy()   # sigmoid(x1)*(1-sigmoid(x1)) => 0.25

高阶梯度

在 tf.GradientTape 上下文管理器中的操作将被记录以进行自动微分。如果在该上下文中计算梯度，则梯度计算也将被记录。因此，完全相同的 API 也适用于高阶梯度。

例如

x = tf.Variable(1.0)  # Create a Tensorflow variable initialized to 1.0

with tf.GradientTape() as t2:
  with tf.GradientTape() as t1:
    y = x * x * x

  # Compute the gradient inside the outer `t2` context manager
  # which means the gradient computation is differentiable as well.
  dy_dx = t1.gradient(y, x)
d2y_dx2 = t2.gradient(dy_dx, x)

print('dy_dx:', dy_dx.numpy())  # 3 * x**2 => 3.0
print('d2y_dx2:', d2y_dx2.numpy())  # 6 * x => 6.0

虽然这确实为您提供了标量函数的二阶导数，但这种模式不会推广到生成 Hessian 矩阵，因为 tf.GradientTape.gradient 仅计算标量的梯度。要构建一个 Hessian 矩阵，请转到 Hessian 示例，该示例位于 雅可比部分 下。

“对 tf.GradientTape.gradient 的嵌套调用”是一种很好的模式，当您从梯度计算标量时，然后生成的标量充当第二次梯度计算的来源，如下面的示例所示。

示例：输入梯度正则化

许多模型容易受到“对抗样本”的影响。这组技术修改模型的输入以混淆模型的输出。最简单的实现（例如 使用快速梯度符号法攻击的对抗样本）沿着输出相对于输入的梯度（“输入梯度”）执行单步操作。

提高对抗样本鲁棒性的一个技术是 输入梯度正则化（Finlay & Oberman，2019），它试图最小化输入梯度的幅度。如果输入梯度很小，则输出的变化也应该很小。

下面是输入梯度正则化的一个朴素实现。该实现是

1. 使用内部磁带计算输出相对于输入的梯度。
2. 计算该输入梯度的幅度。
3. 计算该幅度相对于模型的梯度。

x = tf.random.normal([7, 5])

layer = tf.keras.layers.Dense(10, activation=tf.nn.relu)

with tf.GradientTape() as t2:
  # The inner tape only takes the gradient with respect to the input,
  # not the variables.
  with tf.GradientTape(watch_accessed_variables=False) as t1:
    t1.watch(x)
    y = layer(x)
    out = tf.reduce_sum(layer(x)**2)
  # 1. Calculate the input gradient.
  g1 = t1.gradient(out, x)
  # 2. Calculate the magnitude of the input gradient.
  g1_mag = tf.norm(g1)

# 3. Calculate the gradient of the magnitude with respect to the model.
dg1_mag = t2.gradient(g1_mag, layer.trainable_variables)

[var.shape for var in dg1_mag]

雅可比矩阵

所有前面的示例都对某个源张量（或张量集）相对于标量目标的梯度进行了计算。

雅可比矩阵 表示向量值函数的梯度。每行包含向量中一个元素的梯度。

tf.GradientTape.jacobian 方法允许您有效地计算雅可比矩阵。

请注意

• 与 gradient 相似：sources 参数可以是张量或张量容器。
• 与 gradient 不同：target 张量必须是单个张量。

标量源

作为第一个示例，以下是向量目标相对于标量源的雅可比矩阵。

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

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

dy_dx = tape.jacobian(y, delta)

当您对标量进行雅可比计算时，结果将具有目标的形状，并给出每个元素相对于源的梯度

print(y.shape)
print(dy_dx.shape)

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

张量源

无论输入是标量还是张量，tf.GradientTape.jacobian 都能有效地计算每个源元素相对于每个目标元素（或目标集）的梯度。

例如，此层的输出形状为 (10, 7)

x = tf.random.normal([7, 5])
layer = tf.keras.layers.Dense(10, activation=tf.nn.relu)

with tf.GradientTape(persistent=True) as tape:
  y = layer(x)

y.shape

并且该层的内核形状为 (5, 10)

layer.kernel.shape

输出相对于内核的雅可比矩阵的形状是这两个形状连接在一起

j = tape.jacobian(y, layer.kernel)
j.shape

如果您对目标的维度求和，您将得到由 tf.GradientTape.gradient 计算的总和的梯度

g = tape.gradient(y, layer.kernel)
print('g.shape:', g.shape)

j_sum = tf.reduce_sum(j, axis=[0, 1])
delta = tf.reduce_max(abs(g - j_sum)).numpy()
assert delta < 1e-3
print('delta:', delta)

示例：Hessian

虽然 tf.GradientTape 没有提供用于构建 Hessian 矩阵 的显式方法，但可以使用 tf.GradientTape.jacobian 方法构建一个。

x = tf.random.normal([7, 5])
layer1 = tf.keras.layers.Dense(8, activation=tf.nn.relu)
layer2 = tf.keras.layers.Dense(6, activation=tf.nn.relu)

with tf.GradientTape() as t2:
  with tf.GradientTape() as t1:
    x = layer1(x)
    x = layer2(x)
    loss = tf.reduce_mean(x**2)

  g = t1.gradient(loss, layer1.kernel)

h = t2.jacobian(g, layer1.kernel)

print(f'layer.kernel.shape: {layer1.kernel.shape}')
print(f'h.shape: {h.shape}')

要将此 Hessian 用于 牛顿法 步，您首先需要将其轴展平为矩阵，并将梯度展平为向量

n_params = tf.reduce_prod(layer1.kernel.shape)

g_vec = tf.reshape(g, [n_params, 1])
h_mat = tf.reshape(h, [n_params, n_params])

Hessian 矩阵应该是对称的

def imshow_zero_center(image, **kwargs):
  lim = tf.reduce_max(abs(image))
  plt.imshow(image, vmin=-lim, vmax=lim, cmap='seismic', **kwargs)
  plt.colorbar()

imshow_zero_center(h_mat)

牛顿法更新步骤如下所示

eps = 1e-3
eye_eps = tf.eye(h_mat.shape[0])*eps

# X(k+1) = X(k) - (∇²f(X(k)))^-1 @ ∇f(X(k))
# h_mat = ∇²f(X(k))
# g_vec = ∇f(X(k))
update = tf.linalg.solve(h_mat + eye_eps, g_vec)

# Reshape the update and apply it to the variable.
_ = layer1.kernel.assign_sub(tf.reshape(update, layer1.kernel.shape))

虽然对于单个 tf.Variable 来说这相对简单，但将其应用于非平凡模型将需要仔细的连接和切片才能在多个变量上生成完整的 Hessian。

批处理雅可比矩阵

在某些情况下，您希望对目标堆栈中的每个目标相对于源堆栈中的每个源进行雅可比计算，其中每个目标-源对的雅可比矩阵是独立的。

例如，这里输入 x 的形状为 (batch, ins)，输出 y 的形状为 (batch, outs)

x = tf.random.normal([7, 5])

layer1 = tf.keras.layers.Dense(8, activation=tf.nn.elu)
layer2 = tf.keras.layers.Dense(6, activation=tf.nn.elu)

with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
  tape.watch(x)
  y = layer1(x)
  y = layer2(y)

y.shape

y 相对于 x 的完整雅可比矩阵的形状为 (batch, ins, batch, outs)，即使您只需要 (batch, ins, outs)

j = tape.jacobian(y, x)
j.shape

如果堆栈中每个项目的梯度是独立的，则此张量的每个 (batch, batch) 切片都是一个对角矩阵

imshow_zero_center(j[:, 0, :, 0])
_ = plt.title('A (batch, batch) slice')

def plot_as_patches(j):
  # Reorder axes so the diagonals will each form a contiguous patch.
  j = tf.transpose(j, [1, 0, 3, 2])
  # Pad in between each patch.
  lim = tf.reduce_max(abs(j))
  j = tf.pad(j, [[0, 0], [1, 1], [0, 0], [1, 1]],
             constant_values=-lim)
  # Reshape to form a single image.
  s = j.shape
  j = tf.reshape(j, [s[0]*s[1], s[2]*s[3]])
  imshow_zero_center(j, extent=[-0.5, s[2]-0.5, s[0]-0.5, -0.5])

plot_as_patches(j)
_ = plt.title('All (batch, batch) slices are diagonal')

要获得所需的结果，您可以对重复的 batch 维度求和，或者使用 tf.einsum 选择对角线

j_sum = tf.reduce_sum(j, axis=2)
print(j_sum.shape)
j_select = tf.einsum('bxby->bxy', j)
print(j_select.shape)

在第一步中不使用额外的维度进行计算会更有效。 tf.GradientTape.batch_jacobian 方法正是这样做的

jb = tape.batch_jacobian(y, x)
jb.shape

error = tf.reduce_max(abs(jb - j_sum))
assert error < 1e-3
print(error.numpy())

x = tf.random.normal([7, 5])

layer1 = tf.keras.layers.Dense(8, activation=tf.nn.elu)
bn = tf.keras.layers.BatchNormalization()
layer2 = tf.keras.layers.Dense(6, activation=tf.nn.elu)

with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
  tape.watch(x)
  y = layer1(x)
  y = bn(y, training=True)
  y = layer2(y)

j = tape.jacobian(y, x)
print(f'j.shape: {j.shape}')

plot_as_patches(j)

_ = plt.title('These slices are not diagonal')
_ = plt.xlabel("Don't use `batch_jacobian`")

在这种情况下，batch_jacobian 仍然会运行并返回具有预期形状的某些内容，但其内容含义不明确

jb = tape.batch_jacobian(y, x)
print(f'jb.shape: {jb.shape}')