Il gradiente Tensorflow restituisce nan o Inf

Question

Sto cercando di implementare un modello WGAN-GP utilizzando tensorflow e keras (per i dati sulle frodi con carta di credito da kaggle ).

Ho seguito principalmente il codice di esempio fornito nel sito Web di keras e molti altri codici di esempio su Internet (ma li ho cambiati dall'immagine ai miei dati), ed è piuttosto semplice.

Ma quando voglio aggiornare il critico, il gradiente di perdita rispetto ai pesi del critico diventa tutto nandopo alcuni batch. E questo fa sì che i pesi del critico diventino nane poi i pesi del generatore diventino nan,... Così tutto diventi nan!

Ho usato tf.debugging.enable_check_numericse ho scoperto che il problema sorge perché a -Infappare nel gradiente dopo alcune iterazioni.

Questo è direttamente correlato al termine della penalità del gradiente nella perdita, perché quando lo rimuovo il problema scompare.

Si prega di notare che lo gpstesso non lo è nan, ma quando ottengo il gradiente della perdita rispetto ai pesi critici ( c_gradsnel codice seguente) contiene -Infe quindi in qualche modo diventa tutto nan.

Ho controllato la matematica e l'architettura di rete per possibili errori (come probabilità di scomparsa del gradiente, ecc.) E ho controllato il mio codice per possibili bug per ore e ore. Ma sono bloccato.

Apprezzerei molto se qualcuno riesce a trovare la radice del problema

Nota: tieni presente che la funzione di output e perdita del critico è leggermente diversa dal documento originale (perché sto cercando di renderlo condizionale) ma ciò non ha nulla a che fare con il problema perché come ho detto prima, l'intero problema scompare quando rimuovo semplicemente il termine di penalità del gradiente

Questo è il mio critico:

critic = keras.Sequential([
        keras.layers.Input(shape=(x_dim,), name='c-input'),
        keras.layers.Dense(64, kernel_initializer=keras.initializers.he_normal(), name='c-hidden-1'),
        keras.layers.LeakyReLU(alpha=0.25, name='c-activation-1'),
        keras.layers.Dense(32, kernel_initializer=keras.initializers.he_normal(), name='c-hidden-2'),
        keras.layers.LeakyReLU(alpha=0.25, name='c-activation-2'),
        keras.layers.Dense(2, activation='tanh', name='c-output')
    ], name='critic')

Questa è la mia funzione di penalità gradiente:

def gradient_penalty(self, batch_size, x_real, x_fake):
    # get the random linear interpolation of real and fake data (x hat)
    alpha = tf.random.uniform([batch_size, 1], 0.0, 1.0)
    x_interpolated = x_real + alpha * (x_fake - x_real)
    with tf.GradientTape() as gp_tape:
        gp_tape.watch(x_interpolated)
        # Get the critic score for this interpolated data
        scores = 0.5 * (self.critic(x_interpolated, training=True) + 1.0)
    # Calculate the gradients w.r.t to this interpolated data
    grads = gp_tape.gradient(scores, x_interpolated)
    # Calculate the norm of the gradients
    # Gradient penalty enforces the gradient to stay close to 1.0 (1-Lipschitz constraint)
    gp = tf.reduce_mean(tf.square(tf.norm(grads, axis=-1) - 1.0))
    return gp

E questo è il codice di aggiornamento del critico

# Get random samples from latent space
z = GAN.random_samples((batch_size, self.latent_dim))

# Augment random samples with the class label (1 for class "fraud") for conditioning
z_conditioned = tf.concat([z, tf.ones((batch_size, 1))], axis=1)
# Generate fake data using random samples
x_fake = self.generator(z_conditioned, training=True)

# Calculate the loss and back-propagate
with tf.GradientTape() as c_tape:
    c_tape.watch(x_fake)
    c_tape.watch(x_real)

    # Get the scores for the fake data
    output_fake = 0.5 * (self.critic(x_fake) + 1.0)
    score_fake = tf.reduce_mean(tf.reduce_sum(output_fake, axis=1))
    # Get the scores for the real data
    output_real = 0.5 * (self.critic(x_real, training=True) + 1.0)
    score_real = tf.reduce_mean((1.0 - 2.0 * y_real) * (output_real[:, 0] - output_real[:, 1]))

# Calculate the gradient penalty
gp = self.gp_coeff * self.gradient_penalty(batch_size, x_real, x_fake)
# Calculate critic's loss (added 1.0 so its ideal value becomes zero)
c_loss = 1.0 + score_fake - score_real + gp
# Calculate the gradients
c_grads = c_tape.gradient(c_loss, self.critic.trainable_weights)
# back-propagate the loss
self.c_optimizer.apply_gradients(zip(c_grads, self.critic.trainable_weights))

Nota anche: come puoi vedere, non utilizzo alcuna entropia incrociata o altre funzioni autoscritte con il rischio di divisione per zero.

Answer 1

Quindi, dopo aver scavato molto di più in Internet, si scopre che ciò è dovuto all'instabilità numerica di tf.norm(e anche ad alcune altre funzioni).

Nel caso della normfunzione, il problema è che quando si calcola il suo gradiente, il suo valore appare al denominatore. Quindi d(norm(x))/dxa x = 0diventerebbe 0 / 0(questo è il misterioso division-by-zeroche stavo cercando!)

Il problema è che il grafo computazionale a volte finisce con cose come a / adove a = 0che numericamente non è definito ma il limite esiste. E a causa del modo in cui tensorflow funziona (che calcola i gradienti usando la regola della catena) risulta in nans o +/-Infs.

Il modo migliore probabilmente sarebbe che tensorflow rilevi questi modelli e li sostituisca con il loro equivalente semplificato analiticamente. Ma fino a quando non lo fanno, abbiamo un altro modo, e cioè usare qualcosa chiamato tf.custom_gradientper definire la nostra funzione personalizzata con il nostro gradiente personalizzato ( problema correlato sul loro github )

Anche se nel mio caso c'era in realtà una soluzione ancora più semplice (anche se non era semplice quando non sapevo che tf.normera il colpevole):

Quindi invece di:

tf.norm(x)

Puoi usare:

tf.sqrt(tf.reduce_sum(tf.square(x)) + 1.0e-12)

Nota: fai attenzione alle dimensioni (se x è una matrice o un tensore e devi calcolare le norme per riga o colonna)! questo è solo un codice di esempio per dimostrare il concetto

Spero che aiuti qualcuno