advanced_autodiff.ipynb

Kernel: Python 3

Copyright 2020 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Diferenciação automática avançada

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

O guia Introdução aos gradientes e diferenciação automática inclui tudo o que é necessário para calcular gradientes no TensorFlow. Este guia se concentra em recursos mais profundos e menos comuns da API tf.GradientTape.

Configuração

import tensorflow as tf

import matplotlib as mpl
import matplotlib.pyplot as plt

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

Controlando a gravação de gradientes

No guia de diferenciação automática você viu como controlar quais variáveis ​​e tensores são observados pela fita (tape) durante a construção do cálculo do gradiente.

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

Pare de gravar

Se você quiser interromper a gravação de gradientes, pode usar tf.GradientTape.stop_recording para suspender temporariamente a gravação.

Isto pode ser útil para reduzir a sobrecarga se você não quiser diferenciar uma operação complicada no meio do seu modelo. Isso pode incluir o cálculo de uma métrica ou de 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'])

Reinicie/inicie a gravação do zero

Se você deseja recomeçar do zero, use tf.GradientTape.reset. Simplesmente sair do bloco de fita de gradiente e reiniciar geralmente é mais fácil de ler, mas você pode usar o método reset quando sair do bloco de fita for 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'])

Pare o fluxo do gradiente com precisão

Em contraste com os controles globais de fita, mostrados acima, a função tf.stop_gradient é muito mais precisa. Ela pode ser usada para impedir que gradientes fluam ao longo de um caminho específico, 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'])

Gradientes personalizados

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

1. Não há um gradiente definido para um novo op que você está escrevendo.

2. Os cálculos padrão são numericamente instáveis.

3. Você quer armazenar em cache uma computação cara do passo para frente.

4. Você quer modificar um valor (por exemplo, usando tf.clip_by_value ou tf.math.round) sem modificar o gradiente.

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

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

Aqui está um exemplo que 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

Consulte a documentação da API do decorador tf.custom_gradient para mais detalhes.

Gradientes personalizados em SavedModel

Observação: esse recurso está disponível no TensorFlow 2.6.

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

Para ser salva no SavedModel, a função gradiente deve ser rastreável (para saber mais, confira o guia Melhor desempenho com 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))

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

Múltiplas fitas

Múltiplas 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

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

Gradientes de ordem superior

As operações dentro do gerenciador de contexto tf.GradientTape são gravadas para diferenciação automática. Se forem computados gradientes nesse contexto, a computação do gradiente também será gravada. Como resultado, a mesma API funcionará também para gradientes de ordem superior.

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

Embora isto forneça a segunda derivada de uma função escalar, esse padrão não pode ser generalizado para produzir uma matriz Hessiana, já que tf.GradientTape.gradient calcula apenas o gradiente de um escalar. Para construir uma matriz Hessiana, veja o exemplo Hessiano na seção Jacobiana.

"Chamadas aninhadas para tf.GradientTape.gradient" é um bom padrão quando você está calculando um escalar a partir de um gradiente e, em seguida, o escalar resultante atua como origem para um segundo cálculo de 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 sua saída. A implementação mais simples - como o Exemplo Adversário usando o ataque Fast Gradient Signed Method - dá um único passo ao longo do gradiente da saída em relação à entrada; o "gradiente de entrada".

Uma técnica usada para aumentar a robustez para exemplos adversários é a regularização do gradiente de entrada (Finlay & Oberman, 2019), que tenta minimizar a magnitude do gradiente de entrada. Se o gradiente de entrada for pequeno, a mudança na saída também deverá ser pequena.

Abaixo está uma implementação ingênua da regularização do gradiente de entrada. A implementação consiste em:

1. Calcular o gradiente da saída em relação à entrada usando uma fita interna.

2. Calcular a magnitude desse gradiente de entrada.

3. Calcular 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]

Jacobianos

Todos os exemplos anteriores obtiveram os gradientes de um destino escalar em relação a algum(s) tensor(es) de origem.

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

O método tf.GradientTape.jacobian permite calcular com eficiência uma matriz Jacobiana.

Observe que:

• Como gradient: o argumento sources pode ser um tensor ou um container de tensores.

• Diferentemente de gradient: o tensor target deve ser um único tensor.

Origem escalar

Como primeiro exemplo, eis o Jacobiano de um destino vetorial em relação a uma origem 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ê considera o Jacobiano em relação a um escalar, o resultado tem o formato do destino e fornece o gradiente de cada elemento em relação à origem:

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

Origem do tensor

Independente da entrada ser escalar ou tensor, tf.GradientTape.jacobian calcula com eficiência o gradiente de cada elemento da origem em relação a cada elemento do(s) destino(s).

Por exemplo, a saída desta camada tem o formato (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

E o formato do kernel da camada é (5, 10):

layer.kernel.shape

O formato do Jacobiano da saída em relação ao kernel são esses dois formatos concatenados:

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

Se você somar as dimensões do destino, ficará 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)

Exemplo: Hessiano

Embora tf.GradientTape não forneça um método explícito para construir uma matriz Hessiana, é possível construir uma usando o método tf.GradientTape.jacobian.

Observação: A matriz Hessiana contém N**2 parâmetros. Por esta e outras razões não é útil para a maioria dos modelos. Este exemplo foi incluído mais como uma demonstração de como usar o método tf.GradientTape.jacobian e não é uma recomendação à otimização direta baseada no Hessiano. Um produto vetorial Hessiano pode ser calculado eficientemente com fitas aninhadas e é uma abordagem muito mais eficiente para otimização de segunda ordem.

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

Para usar este Hessiano para um passo do método de Newton, você primeiro deve achatar seus eixos para produzir uma matriz e achatar o gradiente para produzir 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)

O passo de atualização do método de Newton é mostrado abaixo:

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

Observação: Não inverta a matriz.

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

Embora isso seja relativamente simples para uma única tf.Variable, aplicar isto a um modelo não trivial exigiria concatenação e fatiamento cuidadosos para produzir um Hessiano completo em múltiplas variáveis.

Jacobiano em lote

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

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

O Jacobiano completo de y em relação a x tem o formato (batch, ins, batch, outs), mesmo que você queira apenas (batch, ins, outs):

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

Se os gradientes de cada item na pilha forem independentes, então cada (batch, batch) deste tensor é uma matriz diagonal:

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

Para obter o resultado desejado, você pode somar a dimensão batch duplicada 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)

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

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

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

Cuidado: tf.GradientTape.batch_jacobian verifica apenas se a primeira dimensão da origem e do destino correspondem. Não verifica se os gradientes são realmente independentes. Você é quem decide usar batch_jacobian apenas onde fizer sentido. Por exemplo, adicionar tf.keras.layers.BatchNormalization destrói a independência, pois normaliza em toda a dimensão batch:

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`")

Nesse caso, batch_jacobian ainda roda e retorna algo com o formato esperado, mas seu conteúdo tem um significado pouco claro:

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