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

Question

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

但是，如果我的张量中有 NaN 值并且我想在计算总和时简单地忽略这些 NaN 怎么办？有人知道怎么做吗？

附言。我知道我们可以使用 tf.boolean_mask 来过滤 NaN，但如果我只是做 tensor = tf.boolean_mask(tensor, tf.logical_not(tf.is_nan(tensor))，输出将被压缩成一个维度，这不是我想要的。

非常感谢!

Answer 1

您可以使用 tf.where()在保留原始形状的同时用零替换 tensor 中的 NaN 值:

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)