El gradiente de Tensorflow devuelve nan o Inf

Question

Estoy tratando de implementar un modelo WGAN-GP usando tensorflow y keras (para datos de fraude de tarjetas de crédito de kaggle ).

Principalmente seguí el código de muestra que se proporciona en el sitio web de keras y varios otros códigos de muestra en Internet (pero los cambié de imagen a mis datos), y es bastante sencillo.

Pero cuando quiero actualizar el crítico, el gradiente de pérdida de los pesos del crítico se convierte en todo nandespués de unos pocos lotes. Y esto hace que los pesos del crítico se vuelvan nany luego los pesos del generador se vuelvan nan,... ¡Así que todo se vuelve nan!

Utilicé tf.debugging.enable_check_numericsy descubrí que el problema surge porque -Infaparece en el gradiente después de algunas iteraciones.

Esto está directamente relacionado con el término de penalización de gradiente en la pérdida, porque cuando elimino eso, el problema desaparece.

Tenga en cuenta que el gpmismo no lo es nan, pero cuando obtengo el gradiente de los pesos del crítico de pérdida wrt ( c_gradsen el código a continuación) contiene -Infy luego de alguna manera se convierte en todo nan.

Revisé la arquitectura matemática y de red en busca de posibles errores (como la probabilidad de que el gradiente desaparezca, etc.) y revisé mi código en busca de posibles errores durante horas y horas. Pero estoy atascado.

Agradecería mucho si alguien puede encontrar la raíz del problema.

Nota: tenga en cuenta que la función de salida y pérdida del crítico es ligeramente diferente del artículo original (porque estoy tratando de hacerlo condicional), pero eso no tiene nada que ver con el problema porque, como dije antes, todo el problema desaparece. cuando simplemente elimino el término de penalización de gradiente

Esta es mi crítica:

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

Esta es mi función de penalización de 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

Y este es el código de actualización del crítico.

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

También tenga en cuenta: como puede ver, no uso ninguna entropía cruzada u otras funciones autoescritas con el riesgo de división por cero.

Answer 1

Entonces, después de investigar mucho más en Internet, resulta que esto se debe a la inestabilidad numérica de tf.norm(y algunas otras funciones también).

En el caso de la normfunción, el problema es que al calcular su gradiente, su valor aparece en el denominador. Entonces d(norm(x))/dxat x = 0se convertiría 0 / 0(¡este es el misterio division-by-zeroque estaba buscando!)

El problema es que el gráfico computacional a veces termina con cosas como a / aque a = 0numéricamente no está definido pero el límite existe. Y debido a la forma en que funciona tensorflow (que calcula los gradientes usando la regla de la cadena) da como resultado nans o +/-Infs.

Probablemente, la mejor manera sería que tensorflow detectara estos patrones y los reemplazara con su equivalente simplificado analíticamente. Pero hasta que lo hagan, tenemos otra forma, y ​​es usar algo llamado tf.custom_gradientpara definir nuestra función personalizada con nuestro gradiente personalizado ( problema relacionado en su github )

Aunque en mi caso en realidad había una solución aún más simple (aunque no era simple cuando no sabía que tf.normel culpable era él):

Así que en lugar de:

tf.norm(x)

Puedes usar:

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

Nota: ¡Tenga cuidado con las dimensiones (si x es una matriz o un tensor y necesita calcular las normas por filas o columnas)! este es solo un código de muestra para demostrar el concepto

Espero que ayude a alguien