Расширенная автоматическая дифференциация

Посмотреть на TensorFlow.org Запускаем в Google Colab Посмотреть исходный код на GitHub Скачать блокнот

Введение в градиенты и автоматическое дифференцирование руководства включает в себя все необходимое для вычисления градиентов в TensorFlow. Это руководство фокусируется на более глубокие, менее общих черт tf.GradientTape API.

Настраивать

pip uninstall tensorflow keras -y
pip install tf-nightly
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'])
dz/dx: tf.Tensor(4.0, shape=(), dtype=float32)
dz/dy: None
2021-07-01 01:22:12.311927: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-07-01 01:22:12.319895: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-07-01 01:22:12.320536: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-07-01 01:22:12.322087: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2021-07-01 01:22:12.322666: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-07-01 01:22:12.323332: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-07-01 01:22:12.323939: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-07-01 01:22:12.907440: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-07-01 01:22:12.908098: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-07-01 01:22:12.908676: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-07-01 01:22:12.909259: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1510] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14646 MB memory:  -> device: 0, name: NVIDIA Tesla V100-SXM2-16GB, pci bus id: 0000:00:05.0, compute capability: 7.0

Сбросить / начать запись с нуля

Если вы хотите , чтобы начать полностью использовать 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'])
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

Пользовательские градиенты

В некоторых случаях вы можете захотеть точно контролировать, как вычисляются градиенты, а не использовать значение по умолчанию. Эти ситуации включают:

  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.Tensor(2.0, shape=(), dtype=float32)

Обратитесь к tf.custom_gradient документации декоратор API для более подробной информации.

Пользовательские градиенты в SavedModel

Пользовательские градиенты могут быть сохранены в SavedModel с помощью опции tf.saved_model.SaveOptions(experimental_custom_gradients=True) .

Чтобы спастись в 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))
INFO:tensorflow:Assets written to: saved_model/assets
tf.Tensor(2.0, shape=(), dtype=float32)
2021-07-01 01:22:13.395687: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:185] None of the MLIR Optimization Passes are enabled (registered 2)

Примечание о приведенном выше примере: Если вы пытаетесь заменить приведенный выше код с 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
1.0
tape1.gradient(ys, x1).numpy()   # sigmoid(x1)*(1-sigmoid(x1)) => 0.25
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
dy_dx: 3.0
d2y_dx2: 6.0

Несмотря на то , что дает вам вторую производную скалярной функции, эта схема не обобщается для получения матрицы Гессе, так как tf.GradientTape.gradient только вычисляет градиент скалярной. Для того, чтобы построить матрицу Гесса , перейти к Гессу , например под якобиевой секцией .

«Уплотненные вызовы tf.GradientTape.gradient » хороший шаблон , когда вы вычисление скаляра из градиента, а затем полученные скалярный выступают в качестве источника для второго расчета градиента, как показан в следующем примере.

Пример: регуляризация входного градиента

Многие модели восприимчивы к «примерам состязательности». Этот набор методов изменяет входные данные модели, чтобы запутать выходные данные модели. Простейшей реализацией, такие как состязательный например , используя Fast Градиент Signed метод атаки -takes один шаг вдоль градиента выходного сигнала по отношению к входному; «входной градиент».

Один из методов , чтобы увеличить устойчивость к состязательным примерам ввода градиента регуляризация (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]
[TensorShape([5, 10]), TensorShape([10])]

Якобианцы

Во всех предыдущих примерах использовались градиенты скалярной цели относительно некоторого исходного тензора (ов).

Матрица Якоби представляет градиенты вектора функции. Каждая строка содержит градиент одного из элементов вектора.

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)
(201,)
(201,)
plt.plot(x.numpy(), y, label='y')
plt.plot(x.numpy(), dy_dx, label='dy/dx')
plt.legend()
_ = plt.xlabel('x')

PNG

Тензорный источник

Если вход скаляр или тензор, 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
TensorShape([7, 10])

И форма ядра в слое составляет (5, 10) :

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

Форма якобиана вывода по отношению к ядру - это эти две формы, сцепленные вместе:

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

Если просуммировать размеры целевого объекта, вы остаетесь с градиентом от суммы , которая была бы вычисленным 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)
g.shape: (5, 10)
delta: 2.3841858e-07

Пример: гессен

Хотя tf.GradientTape не дает явный способ построения матрицы Гесса , что можно построить один используя 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}')
layer.kernel.shape: (5, 8)
h.shape: (5, 8, 5, 8)

Для того, чтобы использовать этот Гесс для метода Ньютона шага, вы бы сначала выравнивать свои топоры в матрицу, и расплющить градиент в вектор:

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) - (\u2207\xb2f(X(k)))^-1 @ \u2207f(X(k))
# h_mat = \u2207\xb2f(X(k))
# g_vec = \u2207f(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))
2021-07-01 01:22:14.866543: I tensorflow/core/util/cuda_solvers.cc:180] Creating CudaSolver handles for stream 0x7278d90

Хотя это относительно просто для одного tf.Variable , применяя это к нетривиальной модели потребует тщательной конкатенацию и нарезки , чтобы произвести полный Гессе по нескольким переменным.

Пакетный якобиан

В некоторых случаях вы хотите взять якобиан каждой из стека целей по отношению к стеку источников, где якобианы для каждой пары цель-источник независимы.

Например, здесь вход 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)

Было бы гораздо эффективнее провести расчет без дополнительного измерения. tf.GradientTape.batch_jacobian метод делает именно то, что:

jb = tape.batch_jacobian(y, x)
jb.shape
WARNING:tensorflow:5 out of the last 5 calls to <function pfor.<locals>.f at 0x7fb5e8133560> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
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}')
WARNING:tensorflow:6 out of the last 6 calls to <function pfor.<locals>.f at 0x7fb5dc652830> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
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)