Différenciation automatique avancée

Voir sur TensorFlow.org Exécuter dans Google Colab Voir la source sur GitHub Télécharger le cahier

L' introduction à des gradients et la différenciation automatique guide comprend tous les éléments nécessaires à des gradients de calculate dans tensorflow. Ce guide met l' accent sur les caractéristiques plus profondes, moins courantes de la tf.GradientTape API.

Installer

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)

Contrôle de l'enregistrement du dégradé

Dans le guide de différenciation automatique que vous avez vu comment contrôler les variables et tenseurs sont surveillés par la bande tout en construisant le calcul du gradient.

La bande a également des méthodes pour manipuler l'enregistrement.

Arrête d'enregistrer

Si vous souhaitez arrêter l' enregistrement des gradients, vous pouvez utiliser tf.GradientTape.stop_recording de suspendre temporairement l' enregistrement.

Cela peut être utile pour réduire les frais généraux si vous ne souhaitez pas différencier une opération compliquée au milieu de votre modèle. Cela peut inclure le calcul d'une métrique ou d'un résultat intermédiaire :

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

Réinitialiser/démarrer l'enregistrement à partir de zéro

Si vous souhaitez recommencer entièrement, utilisez tf.GradientTape.reset . Sortir tout simplement le bloc de bande de gradient et le redémarrage est généralement plus facile à lire, mais vous pouvez utiliser la reset méthode lors de la sortie du bloc de bande est difficile , voire impossible.

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

Arrêtez le flux de gradient avec précision

Contrairement aux contrôles de bande globaux ci - dessus, la tf.stop_gradient fonction est beaucoup plus précis. Il peut être utilisé pour empêcher les dégradés de s'écouler le long d'un chemin particulier, sans avoir besoin d'accéder à la bande elle-même :

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

Dégradés personnalisés

Dans certains cas, vous souhaiterez peut-être contrôler exactement comment les dégradés sont calculés plutôt que d'utiliser la valeur par défaut. Ces situations incluent :

  1. Il n'y a pas de dégradé défini pour une nouvelle opération que vous écrivez.
  2. Les calculs par défaut sont numériquement instables.
  3. Vous souhaitez mettre en cache un calcul coûteux de la passe avant.
  4. Vous souhaitez modifier une valeur (par exemple, en utilisant tf.clip_by_value ou tf.math.round ) sans modifier le gradient.

Pour le premier cas, d'écrire une nouvelle op , vous pouvez utiliser tf.RegisterGradient pour configurer votre propre (voir l'API docs pour plus de détails). (Notez que le registre de dégradé est global, modifiez-le donc avec prudence.)

Pour les trois derniers cas, vous pouvez utiliser tf.custom_gradient .

Voici un exemple qui applique tf.clip_by_norm au gradient intermédiaire:

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

Reportez - vous aux tf.custom_gradient API décorateur docs pour plus de détails.

Dégradés personnalisés dans SavedModel

Gradients personnalisés peuvent être enregistrés SavedModel en utilisant l'option tf.saved_model.SaveOptions(experimental_custom_gradients=True) .

Pour être sauvé dans le SavedModel, la fonction de gradient doit être traçable (pour en savoir plus, consultez la Meilleure performance avec tf.function guide).

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)

Une note sur l'exemple ci - dessus: Si vous essayez de remplacer le code ci - dessus avec tf.saved_model.SaveOptions(experimental_custom_gradients=False) , le gradient produira toujours le même résultat sur le chargement. La raison en est que le registre de gradient contient encore le gradient personnalisé utilisé dans la fonction call_custom_op . Toutefois, si vous redémarrez le moteur d' exécution après l' enregistrement sans dégradés personnalisés, l' exécution du modèle chargé sous la tf.GradientTape jetteront l'erreur: LookupError: No gradient defined for operation 'IdentityN' (op type: IdentityN) l' LookupError: No gradient defined for operation 'IdentityN' (op type: IdentityN) .

Plusieurs bandes

Plusieurs bandes interagissent de manière transparente.

Par exemple, ici chaque bande regarde un ensemble différent de tenseurs :

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

Dégradés d'ordre supérieur

Opérations à l' intérieur du tf.GradientTape gestionnaire de contexte sont enregistrés pour la différenciation automatique. Si les gradients sont calculés dans ce contexte, le calcul du gradient est également enregistré. En conséquence, la même API fonctionne également pour les gradients d'ordre supérieur.

Par example:

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

Bien que cela ne vous donne la dérivée seconde d'une fonction scalaire, ce modèle ne généralise pas pour produire une matrice hessienne, puisque tf.GradientTape.gradient ne calcule que le gradient d'un scalaire. Pour construire une matrice hessienne , allez à l' exemple hessois sous la section jacobienne .

« Imbrications à tf.GradientTape.gradient » est un bon modèle lorsque vous calculez un scalaire d'un gradient, puis le scalaire résultant agit comme une source pour un second calcul de gradient, comme dans l'exemple suivant.

Exemple : régularisation du gradient d'entrée

De nombreux modèles sont susceptibles d'« exemples contradictoires ». Cette collection de techniques modifie l'entrée du modèle pour confondre la sortie du modèle. La mise en oeuvre plus simple, tel que l' exemple Adversarial en utilisant le gradient rapide Signé attaque Méthode -Prend une seule étape le long du gradient de la sortie par rapport à l'entrée; le "gradient d'entrée".

Une technique pour augmenter la robustesse à des exemples contradictoires est régularisation gradient entrée (Finlay & Oberman, 2019), qui tente de minimiser l'ampleur du gradient d'entrée. Si le gradient d'entrée est petit, alors le changement dans la sortie devrait être petit aussi.

Vous trouverez ci-dessous une implémentation naïve de la régularisation du gradient d'entrée. La mise en œuvre est :

  1. Calculez le gradient de la sortie par rapport à l'entrée à l'aide d'une bande intérieure.
  2. Calculez l'amplitude de ce gradient d'entrée.
  3. Calculez le gradient de cette grandeur par rapport au modèle.
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])]

Jacobiens

Tous les exemples précédents ont pris les gradients d'une cible scalaire par rapport à un ou plusieurs tenseurs source.

La matrice jacobienne représente les gradients d'une fonction à valeurs vectorielles. Chaque ligne contient le dégradé d'un des éléments du vecteur.

La tf.GradientTape.jacobian méthode vous permet de calculer efficacement une matrice jacobienne.

Noter que:

  • Comme gradient : Les sources argument peut être un tenseur ou un conteneur de tenseurs.
  • Contrairement à gradient : L' target tenseur doit être un seul tenseur.

Source scalaire

Comme premier exemple, voici le Jacobien d'un vecteur-cible par rapport à une source scalaire.

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)

Lorsque vous prenez le jacobien par rapport à un scalaire le résultat a la forme de la cible, et donne le gradient du chaque élément par rapport à la source:

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

Tenseur source

Si l'entrée est scalaire ou tenseur, tf.GradientTape.jacobian calcule efficacement le gradient de chaque élément de la source par rapport à chaque élément de la cible (S).

Par exemple, la sortie de cette couche présente une forme 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])

Et est la forme du noyau de la couche (5, 10) :

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

La forme du Jacobien de la sortie par rapport au noyau est ces deux formes concaténées ensemble :

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

Si vous somme sur les dimensions de la cible, vous êtes de gauche avec la pente de la somme qui aurait été calculé par 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

Exemple : toile de jute

Alors que tf.GradientTape ne donne pas une méthode explicite pour la construction d' une matrice hessienne il est possible de construire un en utilisant la tf.GradientTape.jacobian méthode.

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)

Pour utiliser cette hessois pour une méthode de Newton l'étape, vous devez d'abord aplatir ses axes dans une matrice, et aplatir le gradient dans un vecteur:

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 matrice de Hesse doit être symétrique :

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

L'étape de mise à jour de la méthode de Newton est illustrée ci-dessous :

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

Bien que ce soit relativement simple pour un seul tf.Variable , l' application à un modèle non trivial nécessiterait concaténation attention et tranchage pour produire une hessois complète à travers plusieurs variables.

Lot Jacobien

Dans certains cas, vous souhaitez prendre le Jacobien de chacune d'une pile de cibles par rapport à une pile de sources, où les Jacobiens pour chaque paire cible-source sont indépendants.

Par exemple, voici l'entrée x est en forme (batch, ins) et la sortie y est mise en forme (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])

La jacobienne complète de y par rapport à x a une forme de (batch, ins, batch, outs) , même si vous voulez seulement (batch, ins, outs) par (batch, ins, outs) :

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

Si les gradients de chaque élément dans la pile sont indépendants, puis tous (batch, batch) tranche de ce tenseur est une matrice diagonale:

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

Pour obtenir le résultat souhaité, vous pouvez résumer sur le double batch dimension, ou bien sélectionner les diagonales en utilisant 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)

Il serait beaucoup plus efficace de faire le calcul sans la dimension supplémentaire en premier lieu. La tf.GradientTape.batch_jacobian méthode fait exactement cela:

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

Dans ce cas, batch_jacobian fonctionne toujours et retourne quelque chose avec la forme attendue, mais le contenu de l'ont une signification claire:

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