此页面由 Cloud Translation API 翻译。
Switch to English

先进的自动微分

在TensorFlow.org上查看 在Google Colab中运行 在GitHub上查看源代码 下载笔记本

自动微分指南包括计算梯度所需的一切。本指南重点介绍tf.GradientTape API的更深层次,较少见的功能。

建立

 import tensorflow as tf

import matplotlib as mpl
import matplotlib.pyplot as plt

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

控制梯度记录

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

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

如果您想停止记录渐变,可以使用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'])
 
dz/dx: tf.Tensor(4.0, shape=(), dtype=float32)
dz/dy: None

如果您希望完全reset()开始,请使用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'])
 
dz/dx: tf.Tensor(4.0, shape=(), dtype=float32)
dz/dy: None

停止渐变

与上面的全局磁带控件相比, 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'])
 
dz/dx: tf.Tensor(4.0, shape=(), dtype=float32)
dz/dy: None

自定义渐变

在某些情况下,您可能希望精确控制渐变的计算方式,而不是使用默认值。这些情况包括:

  • 您正在编写的新操作没有定义的渐变。
  • 默认计算在数值上不稳定。
  • 您希望从正向缓存高速计算。
  • 您想要修改一个值(例如,使用: tf.clip_by_valuetf.math.round )而不修改渐变。

要编写新操作,可以使用tf.RegisterGradient设置自己的操作。有关详情,请参见该页面。 (请注意,渐变注册表是全局的,因此请谨慎更改。)

对于后三种情况,可以使用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.Tensor(2.0, shape=(), dtype=float32)

有关更多详细信息,请参见tf.custom_gradient装饰器。

多个磁带

多个磁带无缝地交互。例如,这里每个磁带监视不同的张量集:

 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
 
1.0
 tape1.gradient(ys, x1).numpy()   # sigmoid(x1)*(1-sigmoid(x1)) => 0.25
 
0.25

高阶渐变

记录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
 
dy_dx: 3.0
d2y_dx2: 6.0

尽管这确实为您提供了标量函数的二阶导数,但是由于GradientTape.gradient仅计算标量的梯度,因此该模式不能推广生成Hessian矩阵。要构造黑森州,请参见“ 雅可比”部分下的黑森州示例

从渐变计算标量时,“嵌套调用GradientTape.gradient ”是一个很好的模式,然后所得的标量将用作第二个渐变计算的源,如以下示例所示。

示例:输入梯度正则化

许多模型容易受到“对抗性例子”的影响。这套技术修改了模型的输入,以混淆模型的输出。 最简单的实现是沿着输出相对于输入的梯度进行一步。 “输入梯度”。

一种提高对抗性示例的鲁棒性的技术是输入梯度正则化 ,它试图使输入梯度的大小最小化。如果输入梯度很小,那么输出的变化也应该很小。

以下是输入梯度正则化的简单实施。实现是:

  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]
 
[TensorShape([5, 10]), TensorShape([10])]

雅各布主义者

前面所有示例均采用标量目标相对于某些源张量的梯度。

雅可比矩阵表示矢量值函数的梯度。每行都包含矢量元素之一的渐变。

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)
 
(201,)
(201,)

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

png

张量源

无论输入是标量还是张量, 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
 
TensorShape([7, 10])

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

 layer.kernel.shape
 
TensorShape([5, 10])

输出相对于内核的雅可比行列的形状是将这两个形状串联在一起:

 j = tape.jacobian(y, layer.kernel)
j.shape
 
TensorShape([7, 10, 5, 10])

如果对目标尺寸进行求和,则剩下的梯度是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)
 
g.shape: (5, 10)
delta: 4.7683716e-07

示例:黑森州

尽管tf.GradientTape没有提供用于构造Hessian矩阵的显式方法,但可以使用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}')
 
layer.kernel.shape: (5, 8)
h.shape: (5, 8, 5, 8)

要将此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])
 

黑森州矩阵应该是对称的:

 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)
 

png

牛顿方法更新步骤如下所示。

 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
 
TensorShape([7, 6])

y相对于x的完整雅可比行列具有(batch, ins, batch, outs)的形状,即使您只想要(batch, ins, outs)

 j = tape.jacobian(y, x)
j.shape
 
TensorShape([7, 6, 7, 5])

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

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

png

 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')
 

png

为了获得理想的结果,您可以对重复的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)
 
(7, 6, 5)
(7, 6, 5)

在没有额外尺寸的情况下进行计算会更加高效。 GradientTape.batch_jacobian方法可以做到这一点。

 jb = tape.batch_jacobian(y, x)
jb.shape
 
TensorShape([7, 6, 5])
 error = tf.reduce_max(abs(jb - j_sum))
assert error < 1e-3
print(error.numpy())
 
0.0

 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}')
 
j.shape: (7, 6, 7, 5)

 plot_as_patches(j)

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

png

在这种情况下batch_jacobian仍然运行,并返回预期的形状的东西 ,但它的内容有不清楚的含义。

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