O gradiente do Tensorflow retorna nan ou Inf

Question

Estou tentando implementar um modelo WGAN-GP usando tensorflow e keras (para dados de fraude de cartão de crédito do kaggle ).

Eu segui principalmente o código de exemplo fornecido no site keras e vários outros códigos de exemplo na Internet (mas os alterei de imagem para meus dados), e é bastante simples.

Mas quando eu quero atualizar o crítico, o gradiente de pesos do crítico perde tudo nandepois de alguns lotes. E isso faz com que os pesos do crítico se tornem nane depois os pesos do gerador se tornem nan,... Então tudo se torna nan!

Eu usei tf.debugging.enable_check_numericse descobri que o problema surge porque a -Infaparece no gradiente após algumas iterações.

Isso está diretamente relacionado ao termo gradiente-penalidade na perda, pois quando eu removo isso o problema desaparece.

Observe que o gppróprio não é nan, mas quando recebo o gradiente dos pesos do crítico de perda wrt ( c_gradsno código abaixo) ele contém -Infe de alguma forma se torna all nan.

Verifiquei a matemática e a arquitetura de rede em busca de possíveis erros (como probabilidade de desaparecimento de gradiente etc.) e verifiquei meu código em busca de possíveis erros por horas e horas. Mas estou preso.

Agradeceria muito se alguém pudesse encontrar a raiz do problema

Nota: Tenha em mente que a função de saída e perda do crítico é um pouco diferente do artigo original (porque estou tentando torná-lo condicional), mas isso não tem nada a ver com o problema porque como eu disse antes, todo o problema desaparece quando eu apenas removo o termo de penalidade de gradiente

Este é o meu crítico:

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 é a minha função de penalidade 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

E este é o código de atualização do 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))

Nota também: Como você pode ver, eu não uso nenhuma entropia cruzada ou outras funções auto-escritas com o risco de divisão por zero.

Answer 1

Então, depois de muito mais pesquisas na internet, verifica-se que isso se deve à instabilidade numérica de tf.norm(e algumas outras funções também).

No caso de normfunção, o problema é que ao calcular seu gradiente, seu valor aparece no denominador. Então d(norm(x))/dxse x = 0tornaria 0 / 0(este é o misterioso division-by-zeroque eu estava procurando!)

O problema é que o gráfico computacional às vezes acaba com coisas como a / aonde a = 0numericamente é indefinido, mas o limite existe. E por causa da maneira como o tensorflow funciona (que calcula os gradientes usando a regra da cadeia), resulta em nans ou +/-Infs.

A melhor maneira provavelmente seria o tensorflow detectar esses padrões e substituí-los por seu equivalente analiticamente simplificado. Mas até que eles façam isso, temos outra maneira, e isso é usar algo chamado tf.custom_gradientpara definir nossa função personalizada com nosso gradiente personalizado ( problema relacionado em seu github )

Embora no meu caso houvesse uma solução ainda mais simples (embora não fosse simples quando eu não sabia que tf.normera o culpado):

Então, em vez de:

tf.norm(x)

Você pode usar:

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

Nota: Tenha cuidado com as dimensões (se x for uma matriz ou tensor e você precisar calcular as normas de linha ou coluna)! este é apenas um código de exemplo para demonstrar o conceito

Espero que ajude alguém