TensorFlow 梯度与 tf.where 在不应该返回 NaN

Question

Below is reproducible code. If you run it, you will see that in the first sess run, the result is nan, whereas the second case gives the correct gradient value of 0.5. But per tf.where and condition specified, they should return the same value. I also simply don't understand why the tf.where function gradient is nan at 1 or -1, which seem to be totally fine input values to me.

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])}))

Thanks for comments!

Answer 1

As explained in the github issue provided by @mikkola, the problem stems from the internal implementation of tf.where . Basically, both alternatives (and their gradient) are computed, and only the correct part is chosen by multiplication of the conditionnal. Alas, if the gradient is inf or nan for the part that is not selected, even when multiplied by 0 you get a nan that eventually propagates to the result.

Since the issue has been filed in May 2016 (that's tensorflow v0.7!) and not patched since, one can safely assume that this won't be anytime soon and start looking for work around.

The easiest way to fix it is to modify your statements so that they always valid and differentiable, even for values that are not meant to be selected.

A general technique would be to clip the input value inside its valid domain. So in your case for example, you could use

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

In your particular case however it would be simpler to just use

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