Tensorflow 梯度 returns nan 或 Inf

Question

我正在尝试使用 tensorflow 和 keras 实现 WGAN-GP 模型（针对 credit card fraud data from kaggle）。

我主要遵循 keras website 中提供的示例代码和互联网上的其他几个示例代码（但将它们从图像更改为我的数据），而且非常简单。

但是当我想更新critic时，损失的梯度w.r.t critic的权重在几批之后变成了全部nan。这导致评论家的权重变为 nan 之后生成器的权重变为 nan,...所以一切都变为 nan!

我用了tf.debugging.enable_check_numerics，发现问题的出现是因为经过一些迭代后渐变中出现了一个-Inf。

这与损失中的梯度惩罚项直接相关，因为当我删除它时，问题就消失了。

请注意 gp 本身不是 nan，但是当我得到损失的梯度时 w.r.t 评论家的权重（c_grads在下面的代码中）它包含 -Inf 然后以某种方式变成全部 nan.

我检查了数学和网络架构是否存在可能的错误（例如梯度消失的概率等），并且我花了好几个小时检查我的代码是否存在可能的错误。但是我卡住了。

如果有人能找到问题的根源，我将不胜感激

注：请记住，评论家的输出和损失函数与原始论文略有不同（因为我试图使其成为条件）但这与问题无关，因为正如我之前所说，当我只需删除梯度惩罚项

这是我的批评家：

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

这是我的梯度惩罚函数：

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

这是评论家的更新码

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

另请注意：如您所见，我没有使用任何交叉熵或其他有被零除风险的自写函数。

Answer 1

所以在深入研究互联网之后，发现这是因为 tf.norm（以及其他一些函数）的数值不稳定性。

在norm函数的情况下，问题是在计算它的梯度时，它的值出现在分母中。所以x = 0处的d(norm(x))/dx会变成0 / 0（这就是我要找的神秘division-by-zero！）

问题是计算图有时会以 a / a 之类的东西结束，其中 a = 0 在数值上未定义但存在限制。由于 tensorflow 的工作方式（使用链式法则计算梯度），它导致 nans 或 +/-Infs.

最好的方法可能是让 tensorflow 检测这些模式并将它们替换为 analytically-simplified 等价物。但在他们这样做之前，我们还有另一种方法，那就是使用一种叫做 tf.custom_gradient to define our custom function with our custom gradient (related issue on their github)

的东西

虽然在我的情况下实际上有一个更简单的解决方案（虽然当我不知道 tf.norm 是罪魁祸首时它并不简单）：

所以代替：

tf.norm(x)

您可以使用：

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

注意：注意维度（如果x是矩阵或张量，你需要计算row-wise或column-wise范数）！这只是演示概念的示例代码

希望对大家有所帮助

Tensorflow 梯度 returns nan 或 Inf

Tensorflow gradient returns nan or Inf

python

deep-learning

keras

tensorflow

generative-adversarial-network