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.

Diferenciación automática avanzada

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver código fuente en GitHub

Descargar bloc de notas

La guía Introducción a los gradientes y la diferenciación automática incluye todo lo necesario para calcular gradientes en TensorFlow. Esta guía se enfoca en características más profundas y menos comunes de la API tf.GradientTape.

Preparación

import tensorflow as tf

import matplotlib as mpl
import matplotlib.pyplot as plt

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

Controlar el registro de gradiente

En la guía de diferenciación automática vio cómo controlar qué variables y tensores son observados por la cinta mientras se construye el cálculo del gradiente.

La cinta también dispone de métodos para manipular la grabación.

Detener la grabación

Si desea detener la grabación de gradientes, puede usar tf.GradientTape.stop_recording para suspender temporalmente la grabación.

Esto puede ser útil para reducir la sobrecarga si no desea diferenciar una operación complicada a mitad de su modelo. Puede incluir calcular una métrica o un resultado intermedio:

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

Reiniciar/empezar a grabar desde cero

Si desea empezar de cero por completo, use tf.GradientTape.reset. Si simplemente sale del bloque de cinta de gradiente y vuelve a empezar, suele ser más fácil de leer, pero puede usar el método reset cuando salir del bloque de cinta sea difícil o imposible.

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

Detener el flujo de gradiente con precisión

A diferencia de los controles globales de cinta anteriores, la función tf.stop_gradient es mucho más precisa. Puede usarse para detener el flujo de gradientes a lo largo de una ruta concreta, sin necesidad de acceder a la propia cinta:

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

En algunos casos, es posible que desee controlar exactamente cómo se calculan los gradientes en lugar de usar el valor por defecto. Estas situaciones incluyen:

1. No hay un gradiente definido para la nueva op que está escribiendo.

2. Los cálculos por defecto son numéricamente inestables.

3. Desea almacenar en caché un cálculo costoso de la pasada hacia delante.

4. Quiere modificar un valor (por ejemplo, usando tf.clip_by_value o tf.math.round) sin modificar el gradiente.

En el primer caso, para escribir una nueva op puede usar tf.RegisterGradient para configurar la suya propia (consulte los documentos de la API para más detalles y tenga en cuenta que el registro de gradientes es global, así que modifíquelo con precaución).

Para los tres últimos casos, puede usar tf.custom_gradient.

Aquí tiene un ejemplo que aplica tf.clip_by_norm al gradiente intermedio:

# 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 la documentación de la API del decorador tf.custom_gradient para obtener más detalles.

Gradientes personalizados en SavedModel

Nota: Esta función está disponible a partir de TensorFlow 2.6.

Los gradientes personalizados pueden guardarse en SavedModel usando la opción tf.saved_model.SaveOptions(experimental_custom_gradients=True).

Para guardarse en el SavedModel, la función de gradiente debe ser trazable (para saber más, consulte la guía Mejor rendimiento con 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))

Una nota sobre el ejemplo anterior: Si intenta reemplazar el código anterior por tf.saved_model.SaveOptions(experimental_custom_gradients=False), el gradiente seguirá produciendo el mismo resultado al cargar. El motivo es que el registro de gradientes aún contiene el gradiente personalizado usado en la función call_custom_op. Sin embargo, si reinicia el runtime después de guardar sin gradientes personalizados, al ejecutar el modelo cargado bajo la función tf.GradientTape se producirá el error: LookupError: No gradient defined for operation 'IdentityN' (op type: IdentityN).

Múltiples cintas

Las cintas múltiples interactúan a la perfección.

Por ejemplo, aquí cada cinta vigila un 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 orden alto

Las operaciones dentro del administrador de contexto tf.GradientTape se registran para la diferenciación automática. Si se calculan gradientes en ese contexto, también se registra el cálculo del gradiente. Como resultado, la misma API funciona también para gradientes de orden alto.

Por ejemplo:

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

Aunque eso le da la segunda derivada de una función escalar, este patrón no se generaliza para producir una matriz hessiana, ya que tf.GradientTape.gradient sólo calcula el gradiente de un escalar. Para construir una matriz hessiana, vaya al ejemplo hessiano en la sección jacobiana.

"Llamadas anidadas a tf.GradientTape.gradient" es un buen patrón cuando se calcula un escalar a partir de un gradiente, y luego el escalar resultante actúa como fuente para un segundo cálculo de gradiente, como en el siguiente ejemplo.

Ejemplo: Regularización del gradiente de entrada

Muchos modelos son susceptibles a los "ejemplos adversariales". Esta colección de técnicas modifica la entrada del modelo para confundir la salida del mismo. La implementación más sencilla (como el Ejemplo adversarial que utiliza el ataque del Método de Gradiente Rápido por Signos) toma un solo paso a lo largo del gradiente de la salida con respecto a la entrada; el "gradiente de entrada".

Una técnica para aumentar la robustez ante ejemplos adversariales es la regularización del gradiente de entrada (Finlay & Oberman, 2019), que intenta minimizar la magnitud del gradiente de entrada. Si el gradiente de entrada es pequeño, entonces el cambio en la salida también debería serlo.

Aquí tiene una implementación ingenua de la regularización por gradiente de entrada. La implementación es:

1. Calcule el gradiente de la salida con respecto a la entrada usando una cinta interior.

2. Calcule la magnitud de ese gradiente de entrada.

3. Calcule el gradiente de esa magnitud con respecto al 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 los ejemplos anteriores tomaban los gradientes de un objetivo escalar con respecto a algún(os) tensor(es) fuente.

La matriz Jacobiana representa los gradientes de una función con valor vectorial. Cada fila contiene el gradiente de uno de los elementos del vector.

El método tf.GradientTape.jacobian permite calcular eficazmente una matriz jacobiana.

Tenga en cuenta que:

• Como gradient: El argumento sources puede ser un tensor o un contenedor de tensores.

• A diferencia de gradient: El tensor target debe ser un único tensor.

Fuente escalar

Como primer ejemplo, he aquí la jacobiana de un vector-objetivo con respecto a un escalar-fuente.

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)

Cuando se toma el jacobiano con respecto a un escalar el resultado tiene la forma del objetivo, y da el gradiente de cada elemento con respecto al origen:

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

Fuente de tensor

Tanto si la entrada es escalar como tensorial, tf.GradientTape.jacobian calcula eficazmente el gradiente de cada elemento de la fuente con respecto a cada elemento del objetivo u objetivos.

Por ejemplo, la salida de esta capa tiene una 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

Y la forma del kernel de la capa es (5, 10):

layer.kernel.shape

La forma del jacobiano de la salida con respecto al kernel son esas dos formas concatenadas:

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

Si suma sobre las dimensiones del objetivo, le queda el gradiente de la suma que habría calculado 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)

Ejemplo: Hessiana

Aunque tf.GradientTape no da un método explícito para construir una matriz hessiana es posible construir una usando el método tf.GradientTape.jacobian.

Nota: La matriz hessiana contiene N**2 parámetros. Por este y otros motivos no es práctico para la mayoría de los modelos. Este ejemplo se incluye más como demostración de cómo usar el método tf.GradientTape.jacobian, y no para apoyar que la optimización directa se base en la hessiana. Un producto vectorial hessiano puede calcularse eficientemente con cintas anidadas, y es un enfoque mucho más eficiente para la optimización de segundo orden.

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 esta hessiana para un paso del método de Newton, primero habría que aplanar sus ejes en una matriz, y aplanar el gradiente en un vector:

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

La matriz hessiana debe 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)

A continuación se muestra el paso de actualización del método de Newton:

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

Nota: En realidad no invierta la 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))

Aunque esto es relativamente sencillo para una sola tf.Variable, aplicarlo a un modelo no trivial requeriría una cuidadosa concatenación y corte para producir un hessiano completo a través de múltiples variables.

Jacobiano de lote

En algunos casos, se desea tomar la jacobiana de cada uno de un apilamiento de objetivos con respecto a un apilamiento de fuentes, donde las jacobianas de cada par objetivo-origen son independientes.

Por ejemplo, aquí la entrada x tiene forma (batch, ins) y la salida y tiene forma (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

La jacobiana completa de y con respecto a x tiene una forma de (batch, ins, batch, outs), aunque sólo quiera (batch, ins, outs):

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

Si los gradientes de cada artículo de la pila son independientes, entonces cada (batch, batch) corte de este tensor es una 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 obtener el resultado deseado, puede sumar sobre la dimensión duplicada batch, o bien seleccionar las diagonales utilizando 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)

Sería mucho más eficiente hacer el cálculo sin la dimensión extra en primer lugar. El método tf.GradientTape.batch_jacobian hace exactamente eso:

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

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

Precaución: tf.GradientTape.batch_jacobian sólo comprueba que la primera dimensión del origen y del objetivo coincidan. No comprueba que los gradientes sean realmente independientes. Depende de usted asegurarse de que sólo usa batch_jacobian cuando tiene sentido. Por ejemplo, añadir un tf.keras.layers.BatchNormalization destruye la independencia, ya que normaliza a través de la dimensión 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`")

En este caso, batch_jacobian aún se ejecuta y devuelve algo con la forma esperada, pero su contenido tiene un significado poco claro:

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