python - tf.where 的 TensorFlow 梯度在不应该返回 NaN 时返回 NaN

Question

下面是可重现的代码。如果运行它，您将看到在第一次 sess 运行中，结果为 nan，而第二种情况给出了正确的梯度值 0.5。但根据指定的 tf.where 和条件，它们应该返回相同的值。我也根本不明白为什么 tf.where 函数梯度在 1 或 -1 时为 nan，这对我来说似乎是完全好的输入值。

tf.reset_default_graph()
x = tf.get_variable('x', shape=[1])
condition = tf.less(x, 0.0)
output = tf.where(condition, -tf.log(-x + 1), tf.log(x + 1))
deriv = tf.gradients(output, x)
with tf.Session() as sess:
    print(sess.run(deriv, {x:np.array([-1])}))

logg = -tf.log(-x+1)
derivv = tf.gradients(logg, x)
with tf.Session() as sess:
    print(sess.run(derivv, {x:np.array([-1])}))

感谢您的评论!

Answer 1

如 github issue 中所述由@mikkola提供，问题源于tf.where的内部实现。基本上，两种选择(及其梯度)都会被计算，并且通过条件的乘法仅选择正确的部分。唉，如果未选择的部分的梯度是inf或nan，即使乘以0你也会得到nan 最终传播到结果。

由于该问题已于 2016 年 5 月提交(即 tensorflow v0.7!)且尚未修补，因此可以放心地假设该问题不会很快出现，并开始寻找解决方法。

修复此问题的最简单方法是修改您的语句，使它们始终有效且可微分，即使对于不打算选择的值也是如此。

通用技术是将输入值限制在其有效域内。例如，在您的情况下，您可以使用

cond = tf.less(x, 0.0)
output = tf.where(cond,
  -tf.log(-tf.where(cond, x, 0) + 1),
  tf.log(tf.where(cond, 0, x) + 1))

但是，在您的具体情况下，使用它会更简单

output = tf.sign(x) * tf.log(tf.abs(x) + 1)