Diferenciação automática avançada

Ver no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

A Introdução ao gradientes e diferenciação automática guia inclui tudo o necessário para gradientes de calcular em TensorFlow. Este guia concentra-se em mais profundas características menos comuns do tf.GradientTape API.

Configurar

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)

Controle de gravação de gradiente

Na guia de diferenciação automática você viu como controlar quais variáveis e tensores são vigiados pela fita ao construir o cálculo do gradiente.

A fita também possui métodos para manipular a gravação.

Pare de gravar

Se você deseja parar a gravação gradientes, você pode usar tf.GradientTape.stop_recording de suspender temporariamente a gravação.

Isso pode ser útil para reduzir a sobrecarga se você não quiser diferenciar uma operação complicada no meio de seu modelo. Isso pode incluir o cálculo de uma métrica ou um resultado intermediário:

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

Reiniciar / iniciar a gravação do zero

Se você deseja começar de novo inteiramente, use tf.GradientTape.reset . Basta sair do bloco de fita gradiente e reiniciar normalmente é mais fácil de ler, mas você pode usar a reset método quando sair do bloco de fita é difícil ou impossível.

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

Pare o fluxo de gradiente com precisão

Em contraste com os controles de fita globais acima, o tf.stop_gradient função é muito mais preciso. Ele pode ser usado para impedir que gradientes fluam ao longo de um determinado caminho, sem a necessidade de acesso à própria fita:

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

Gradientes personalizados

Em alguns casos, você pode querer controlar exatamente como os gradientes são calculados em vez de usar o padrão. Essas situações incluem:

  1. Não há gradiente definido para uma nova operação que você está escrevendo.
  2. Os cálculos padrão são numericamente instáveis.
  3. Você deseja armazenar em cache um cálculo caro do passe para frente.
  4. Pretende modificar um valor (por exemplo, usando tf.clip_by_value ou tf.math.round ) sem modificar o gradiente.

Para o primeiro caso, para escrever um novo op você pode usar tf.RegisterGradient para criar o seu próprio (consulte a documentação da API para mais detalhes). (Observe que o registro de gradiente é global, portanto, altere-o com cuidado.)

Para os três últimos casos, você pode usar tf.custom_gradient .

Aqui é um exemplo que se aplica tf.clip_by_norm ao gradiente intermediário:

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

Consulte as tf.custom_gradient docs decorador API para mais detalhes.

Gradientes personalizados em SavedModel

Gradientes personalizados podem ser salvos em SavedModel usando a opção tf.saved_model.SaveOptions(experimental_custom_gradients=True) .

Para ser salvo no SavedModel, a função de inclinação devem ser rastreáveis (para saber mais, confira o desempenho melhor com tf.function guia).

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)

Uma nota sobre o exemplo acima: se você tentar substituir o código acima com tf.saved_model.SaveOptions(experimental_custom_gradients=False) , o gradiente ainda irá produzir o mesmo resultado no carregamento. A razão é que o registo gradiente ainda contém o gradiente personalizado usado na função call_custom_op . No entanto, se você reiniciar o tempo de execução depois de salvar, sem gradientes personalizados, correndo o modelo carregado sob o tf.GradientTape vai jogar o erro: LookupError: No gradient defined for operation 'IdentityN' (op type: IdentityN) .

Múltiplas fitas

Várias fitas interagem perfeitamente.

Por exemplo, aqui cada fita assiste a um conjunto diferente de tensores:

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

Gradientes de ordem superior

Operações dentro do tf.GradientTape gerente de contexto são registrados para a diferenciação automática. Se os gradientes forem calculados nesse contexto, o cálculo do gradiente também será registrado. Como resultado, a mesma API funciona para gradientes de ordem superior também.

Por exemplo:

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

Enquanto que lhe dá o segundo derivado de uma função escalar, este padrão não generaliza para produzir uma matriz de Hesse, desde tf.GradientTape.gradient única calcula o gradiente de um escalar. Para construir uma matriz de Hesse , ir para o exemplo de Hesse sob a secção Jacobiana .

"Nested chamadas para tf.GradientTape.gradient " é um padrão bom quando você está calculando um escalar de um gradiente, e depois a escalar resultando atua como uma fonte para um segundo cálculo gradiente, como no exemplo a seguir.

Exemplo: regularização de gradiente de entrada

Muitos modelos são suscetíveis a "exemplos adversários". Esta coleção de técnicas modifica a entrada do modelo para confundir a saída do modelo. A mais simples aplicação, tal como o exemplo Adversarial utilizando o gradiente rápida Assinado Método ataque -takes um único passo ao longo do gradiente da saída em relação à entrada; o "gradiente de entrada".

Uma técnica para aumentar a robustez aos exemplos contraditório é entrada gradiente de regularização (Finlay & Oberman, 2019), que as tentativas para minimizar a magnitude do gradiente de entrada. Se o gradiente de entrada for pequeno, a mudança na saída também deve ser pequena.

Abaixo está uma implementação ingênua de regularização de gradiente de entrada. A implementação é:

  1. Calcule o gradiente da saída em relação à entrada usando uma fita interna.
  2. Calcule a magnitude desse gradiente de entrada.
  3. Calcule o gradiente dessa magnitude em relação ao modelo.
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])]

Jacobianos

Todos os exemplos anteriores pegaram os gradientes de um alvo escalar em relação a alguns tensores de origem.

A matriz Jacobiana representa os gradientes de uma função vectorial valorizado. Cada linha contém o gradiente de um dos elementos do vetor.

O tf.GradientTape.jacobian método permite calcular de forma eficiente uma matriz Jacobiana.

Observe que:

  • Como gradient : O sources argumento pode ser um tensor ou um recipiente de tensores.
  • Ao contrário de gradient : O target tensor deve ser um único tensor.

Fonte escalar

Como um primeiro exemplo, aqui está o Jacobiano de um alvo vetorial em relação a uma fonte escalar.

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)

Quando você toma o Jacobian no que diz respeito a um escalar o resultado tem a forma do alvo, e dá o gradiente da cada elemento com relação à fonte:

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

Fonte de tensor

Se a entrada é escalar ou tensor, tf.GradientTape.jacobian calcula de forma eficiente o gradiente de cada elemento de fonte com respeito a cada um dos elementos do alvo (s).

Por exemplo, a saída de esta camada tem uma forma de (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])

E forma do núcleo da camada é (5, 10) :

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

A forma do Jacobiano da saída em relação ao kernel são essas duas formas concatenadas:

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

Se você somar sobre as dimensões do alvo, você é deixado com o gradiente da soma que teria sido calculado por 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

Exemplo: Hessian

Enquanto tf.GradientTape não dá um método explícito para a construção de uma matriz de Hessian é possível construir um usando o tf.GradientTape.jacobian método.

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)

Para utilizar este Hessian por um método de Newton etapa, você primeiro achatar seus eixos em uma matriz, e achatar o gradiente em um vetor:

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

A matriz Hessiana deve ser simétrica:

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

A etapa de atualização do método de Newton é mostrada abaixo:

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

Enquanto isso é relativamente simples para um único tf.Variable , aplicando-se a um modelo não-trivial exigiria concatenação cuidadoso e corte para produzir uma Hessian completo em várias variáveis.

Lote Jacobiano

Em alguns casos, você deseja obter o Jacobiano de cada um de uma pilha de destinos em relação a uma pilha de origens, em que os Jacobianos de cada par destino-origem são independentes.

Por exemplo, aqui a entrada x é moldada (batch, ins) e a saída y é moldada (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])

O Jacobian cheio de y em relação à x tem uma forma de (batch, ins, batch, outs) , mesmo se só deseja (batch, ins, outs) :

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

Se os gradientes de cada item na pilha são independentes, em seguida, cada (batch, batch) fatia de este tensor é uma matriz diagonal:

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

Para obter o resultado desejado, você pode somar sobre o duplicado batch dimensão, ou então selecionar as diagonais usando 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)

Seria muito mais eficiente fazer o cálculo sem a dimensão extra em primeiro lugar. O tf.GradientTape.batch_jacobian método faz exatamente isso:

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

Neste caso, batch_jacobian ainda corre e retorna algo com a forma esperada, mas o seu conteúdo tem um significado claro:

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