Numpy Tensorflow无限遮罩破坏梯度

我正在尝试对选定的索引执行softmax，使用无限遮罩消除不需要的索引。然而，这些不需要的实体的梯度变为

nan

，而不是

0

我没有使用布尔掩码的原因是，在我的批处理中，掩码索引是不同的，这不能得到一个好的矩阵形式。如果这里有解决办法，我会非常乐意采纳

我测试无限掩码的代码是

import numpy as np
import tensorflow as tf

a = tf.placeholder(tf.float32, [5])
inf_mask = tf.placeholder(tf.float32, [5])

b = tf.multiply(a, inf_mask)
sf = tf.nn.softmax(b)

loss = (sf[2] - 0)
grad = tf.gradients(loss, a)

sess = tf.Session()

a_np = np.ones([5])
np_mask = np.ones([5]) * 4
np_mask[1] = -np.inf

print sess.run([sf, grad], feed_dict={
    a: a_np,
    inf_mask: np_mask
})

sess.close()

输出是

[array([ 0.25,  0.  ,  0.25,  0.25,  0.25], dtype=float32), [array([-0.25,   nan,  0.75, -0.25, -0.25], dtype=float32)]]

遮罩正在工作，但梯度有一个

nan

，我想应该是

0

。

如中所述，解决方案是将遮罩定义为由零和

np.inf

组成，并将其相减，而不是相乘。如中所述，解决方案是将掩码定义为由零和

np.inf

组成，并将其相减，而不是相乘。