如何在Tensorflow更新梯度时检查NaN？

Question

大家好，

当您使用大量样本来训练一个大模型时，有些样本可能会导致梯度为NaN，从而导致参数更新失败。

我想找出这些样本，并且同时不希望这批样本的梯度用于更新模型参数，因为这有可能导致模型参数变成NaN。

有没有什么好的方法解决这个问题呢？

我的代码如下：

    # Create an optimizer.
    params = tf.trainable_variables()
    opt = tf.train.AdamOptimizer(1e-3)
    gradients = tf.gradients(self.loss, params)

    max_gradient_norm = 10
    clipped_gradients, self.gradient_norms = tf.clip_by_global_norm(gradients,
                                                     max_gradient_norm)

    self.optimizer = opt.apply_gradients(zip(clipped_gradients, params))

- Issac

Answer 1

您可以通过 tf.check_numerics 检查梯度是否为 NaN：

grad_check = tf.check_numerics(clipped_gradients)
with tf.control_dependencies([grad_check]):
  self.optimizer = opt.apply_gradients(zip(clipped_gradients, params))

grad_check 会在裁剪梯度为 NaN 或无穷大时抛出 InvalidArgument 异常。

tf.control_dependencies 确保在应用梯度之前先评估 grad_check。

还可以查看 tf.add_check_numerics_ops()。

- yuefengz

---

非常感谢。但我又有一个问题。当我在我的代码中添加tf.add_check_numerics_ops()时，会出现内存不足错误。而删除此行，则一切正常。我的模型非常大，所以tf.add_check_numerics_ops()会为检查操作分配更多的GPU内存吗？ - Issac

tf.check_numerics核函数中有一个张量拷贝：https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/check_numerics_op.cc#L60。因此，只需使用 tf.check_numerics 添加一些检查即可。如果想使用 tf.add_check_numerics_ops()，则需要额外注意，因为它会对所有浮点类型的张量运行 tf.check_numerics。 - yuefengz

Answer 2

你可以使用tf.is_nan与tf.cond结合起来，仅在损失不是NaN时执行代码的其余部分。

- Mr_and_Mrs_D