TensorFlow中tf.where梯度不应为NaN却返回NaN的问题排查

首先我先补全你没写完的代码片段（推测是涉及平方根这类在部分输入下会产生NaN的操作），类似这样：

import tensorflow as tf

tf.reset_default_graph()
x = tf.get_variable('x', shape=[1], initializer=tf.constant_initializer(1.0))
condition = tf.less(x, 0.0)
# 假设分支是对x取平方根，负数分支处理为sqrt(-x)的相反数
output = tf.where(condition, -tf.sqrt(-x), tf.sqrt(x))
grad = tf.gradients(output, x)[0]

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print("First run grad:", sess.run(grad))  # 输出NaN
    print("Second run grad:", sess.run(grad)) # 输出0.5

为什么第一次运行会返回NaN？ #

这是TensorFlow 1.x中tf.where的梯度计算机制导致的：

• tf.where在计算梯度时，会同时计算两个分支的梯度张量，再根据条件选择最终的梯度值。
• 当x=1.0时，condition为False，本应只取tf.sqrt(x)的梯度（即0.5），但TensorFlow仍会执行另一个分支-tf.sqrt(-x)的计算——此时-x=-1.0，tf.sqrt(-1.0)会产生NaN，对应的梯度自然也是NaN。
• 虽然最终tf.where会选择正确分支的梯度，但由于另一个分支的梯度张量中存在NaN，在TensorFlow的计算流中，这个NaN会“污染”整个梯度结果，导致第一次运行返回NaN。

至于第二次运行能得到正确结果，是因为TensorFlow会缓存计算图的中间结果，第二次运行时无效分支的NaN计算会被优化跳过，或者直接复用缓存的有效分支梯度，从而得到正确值。

为什么输入1或-1时梯度会是NaN？ #

拿x=1.0举例：无效分支是-tf.sqrt(-x)，输入-1.0给sqrt会产生NaN，其梯度也是NaN；同理，当x=-1.0时，无效分支是tf.sqrt(x)，输入负数给sqrt同样产生NaN，进而导致梯度为NaN。

解决方案 #

要避免这个问题，你需要让TensorFlow只执行符合条件的分支，而不是同时计算两个分支。推荐用tf.cond代替tf.where，它会根据条件动态选择执行分支，完全跳过另一个分支的计算：

import tensorflow as tf

tf.reset_default_graph()
x = tf.get_variable('x', shape=[1], initializer=tf.constant_initializer(1.0))

# 用tf.cond替代tf.where，只执行符合条件的分支
output = tf.cond(
    tf.less(x, 0.0),
    lambda: -tf.sqrt(-x),  # 条件为True时执行
    lambda: tf.sqrt(x)     # 条件为False时执行
)
grad = tf.gradients(output, x)[0]

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print("First run grad:", sess.run(grad))  # 直接输出0.5
    print("Second run grad:", sess.run(grad)) # 输出0.5

如果是在TensorFlow 2.x环境下，这个问题会自然消失——因为TF2的自动梯度（Autograd）是基于执行追踪的，只会计算实际运行的分支的梯度，不需要手动替换tf.where。

内容的提问来源于stack exchange，提问作者StatsNoob