TensorFlow：是否可以在忽略NaN值的同时减少总和？

Question

假设我有一个2-D张量，(batch_size, loss_dim)我希望得到每个数据样本的每个损失维数的总和，可以用来完成tf.reduce_mean(tensor, axis=-1)。

但是，如果我的张量中有NaN值，而我想在计算总和时忽略这些NaN怎么办？有谁知道这是怎么做到的吗？

PS。我知道我们可以tf.boolean_mask用来拟合NaN，但是如果我简单地这样做tensor = tf.boolean_mask(tensor, tf.logical_not(tf.is_nan(tensor))，输出将被压缩为一个维度，这不是我想要的。

非常感谢！

Answer 1

您可以使用tf.where()NaN值替换tensor为零，同时保留原始形状：

tensor = ...

# Replace all NaN values with 0.0.
tensor_without_nans = tf.where(tf.is_nan(tensor), tf.zeros_like(tensor), tensor)

sum_ignoring_nans = tf.reduce_sum(tensor_without_nans, axis=-1)