Differenziazione automatica avanzata

Visualizza su TensorFlow.org Esegui in Google Colab Visualizza la fonte su GitHub Scarica il taccuino

L' Introduzione ai gradienti e automatica differenziazione guida include tutto il necessario per gradienti Calcolare in tensorflow. Questa guida si concentra sulle più profonde, caratteristiche meno comuni del tf.GradientTape API.

Impostare

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)

Controllo della registrazione del gradiente

Nella guida automatica differenziazione si è visto come controllare quali variabili e tensori sono guardati dal nastro, mentre la costruzione il calcolo del gradiente.

Il nastro ha anche metodi per manipolare la registrazione.

Interrompi registrazione

Se si desidera interrompere la registrazione sfumature, è possibile utilizzare tf.GradientTape.stop_recording di sospendere temporaneamente la registrazione.

Questo può essere utile per ridurre il sovraccarico se non si desidera differenziare un'operazione complicata nel mezzo del modello. Ciò potrebbe includere il calcolo di una metrica o di un risultato 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'])
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

Ripristina/avvia la registrazione da zero

Se si desidera ricominciare da capo tutto, utilizzare tf.GradientTape.reset . Basta uscire dal blocco del nastro gradiente e il riavvio di solito è più facile da leggere, ma è possibile utilizzare il reset metodo quando uscendo il blocco del nastro è difficile o impossibile.

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

Arresta il flusso del gradiente con precisione

In contrasto con i comandi del nastro globali sopra, la tf.stop_gradient funzione è molto più precisa. Può essere utilizzato per impedire ai gradienti di scorrere lungo un particolare percorso, senza bisogno di accedere al nastro stesso:

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

Gradienti personalizzati

In alcuni casi, potresti voler controllare esattamente come vengono calcolati i gradienti anziché utilizzare l'impostazione predefinita. Queste situazioni includono:

  1. Non esiste un gradiente definito per una nuova operazione che stai scrivendo.
  2. I calcoli predefiniti sono numericamente instabili.
  3. Si desidera memorizzare nella cache un calcolo costoso dal passaggio in avanti.
  4. Si desidera modificare un valore (ad esempio, utilizzando tf.clip_by_value o tf.math.round ) senza modificare il gradiente.

Per il primo caso, a dare una nuova op è possibile utilizzare tf.RegisterGradient per impostare il proprio (vedere la documentazione API per i dettagli). (Nota che il registro del gradiente è globale, quindi cambialo con cautela.)

Per questi ultimi tre casi, è possibile utilizzare tf.custom_gradient .

Ecco un esempio che applica 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
tf.Tensor(2.0, shape=(), dtype=float32)

Fare riferimento alle tf.custom_gradient documentazione decoratore API per maggiori dettagli.

Gradienti personalizzati in SavedModel

Gradienti personalizzati possono essere salvati su SavedModel utilizzando l'opzione tf.saved_model.SaveOptions(experimental_custom_gradients=True) .

Per essere salvati nella SavedModel, la funzione gradiente deve essere tracciabile (per saperne di più, controllare le prestazioni meglio con tf.function guida).

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)

Una nota riguardo l'esempio precedente: se si tenta di sostituire il codice precedente con tf.saved_model.SaveOptions(experimental_custom_gradients=False) , il gradiente sarà ancora lo stesso risultato sul caricamento. La ragione è che il Registro di gradiente contiene ancora il gradiente personalizzato utilizzato nella funzione call_custom_op . Tuttavia, se si riavvia il runtime dopo il salvataggio, senza sfumature personalizzate, eseguire il modello di carico sotto il tf.GradientTape getterà l'errore: LookupError: No gradient defined for operation 'IdentityN' (op type: IdentityN) .

Più nastri

Più nastri interagiscono perfettamente.

Ad esempio, qui ogni nastro osserva un diverso insieme di tensori:

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

Gradienti di ordine superiore

Le operazioni all'interno del tf.GradientTape contesto responsabile sono registrati per la differenziazione automatica. Se i gradienti vengono calcolati in quel contesto, viene registrato anche il calcolo del gradiente. Di conseguenza, la stessa identica API funziona anche per gradienti di ordine superiore.

Per esempio:

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

Mentre che fa dare la derivata seconda di una funzione scalare, questo schema non generalizza per produrre una matrice Hessiana, poiché tf.GradientTape.gradient calcola solo il gradiente di uno scalare. Per costruire una matrice Hessiana , andare esempio dell'Assia sotto la sezione Jacobiana .

"Nidificati chiama tf.GradientTape.gradient " è un buon modello quando si calcola uno scalare da un gradiente, e quindi lo scalare risultante agisce come una fonte per un secondo calcolo gradiente, come il seguente esempio.

Esempio: regolarizzazione del gradiente di input

Molti modelli sono suscettibili di "esempi contraddittori". Questa raccolta di tecniche modifica l'input del modello per confondere l'output del modello. Il più semplice attuazione, come l' esempio contraddittorio utilizzando il gradiente veloce firmata attacco Metodo -takes un singolo passo lungo la pendenza della curva rispetto all'ingresso; il "gradiente di ingresso".

Una tecnica per aumentare la robustezza agli esempi contraddittorio è ingresso gradiente regolarizzazione (Finlay & Oberman, 2019), che tenta di minimizzare l'entità del gradiente di ingresso. Se il gradiente di input è piccolo, anche il cambiamento nell'output dovrebbe essere piccolo.

Di seguito è riportata un'implementazione ingenua della regolarizzazione del gradiente di input. L'implementazione è:

  1. Calcolare il gradiente dell'output rispetto all'input utilizzando un nastro interno.
  2. Calcola la grandezza di quel gradiente di input.
  3. Calcola il gradiente di quella grandezza rispetto al modello.
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])]

giacobini

Tutti gli esempi precedenti hanno preso i gradienti di un bersaglio scalare rispetto ad alcuni tensori sorgente.

La matrice Jacobiana rappresenta il gradiente di una funzione vettoriale prezioso. Ogni riga contiene il gradiente di uno degli elementi del vettore.

Il tf.GradientTape.jacobian metodo consente di calcolare in modo efficiente una matrice Jacobiana.

Notare che:

  • Come gradient : L' sources argomento può essere un tensore o di un contenitore di tensori.
  • Diversamente gradient : Il target tensore deve essere un unico tensore.

Sorgente scalare

Come primo esempio, ecco lo Jacobiano di un vettore-bersaglio rispetto a una sorgente scalare.

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 si prende la Jacobiana rispetto ad uno scalare il risultato ha la forma del bersaglio, e dà il gradiente della ciascun elemento rispetto alla sorgente:

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

Sorgente tensoriale

Se l'input è scalare o tensore, tf.GradientTape.jacobian calcola efficientemente il gradiente di ciascun elemento della sorgente rispetto a ogni elemento (s) la porta.

Ad esempio, l'uscita di questo strato ha una forma (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 la forma del kernel del livello è (5, 10) :

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

La forma dello Jacobiano dell'output rispetto al kernel sono quelle due forme concatenate insieme:

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

Se si somma su dimensioni del bersaglio, si è lasciato con il gradiente della somma che sarebbe stato determinato da 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

Esempio: Hessian

Mentre tf.GradientTape non dà un metodo esplicito per la costruzione di una matrice Hessiana è possibile costruire uno utilizzando il tf.GradientTape.jacobian metodo.

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)

Per utilizzare questa Hesse per un metodo di Newton passo, si dovrebbe prima appiattirsi suoi assi in una matrice, e appiattire il gradiente in un vettore:

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 hessiana dovrebbe essere simmetrica:

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

La fase di aggiornamento del metodo di Newton è mostrata di seguito:

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

Anche se questo è relativamente semplice per un singolo tf.Variable , l'applicazione di questo a un modello non banale richiederebbe un'attenta concatenazione e affettare a produrre una tela di iuta pieno in più variabili.

Lotto Jacobiano

In alcuni casi, vuoi prendere lo Jacobiano di ciascuno di uno stack di bersagli rispetto a uno stack di sorgenti, dove gli Jacobiani per ciascuna coppia di destinazione-sorgente sono indipendenti.

Ad esempio, qui l'ingresso x è sagomata (batch, ins) e l'uscita y è sagomata (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])

Il Jacobiano piena di y rispetto ad x ha una forma di (batch, ins, batch, outs) , anche se desideri solo (batch, ins, outs) :

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

Se i gradienti di ogni elemento dello stack sono indipendenti, allora ogni (batch, batch) fetta di questo tensore è una 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

Per ottenere il risultato desiderato, è possibile sommare il duplicato batch dimensioni, oppure selezionare le diagonali utilizzando 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)

Sarebbe molto più efficiente eseguire il calcolo senza la dimensione extra in primo luogo. Il tf.GradientTape.batch_jacobian metodo fa esattamente questo:

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

In questo caso, batch_jacobian corre ancora e torna qualcosa con la forma previsto, ma il suo contenuto avere un significato poco chiaro:

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