如何:忽略 NA 值(或以其他方式屏蔽损失)的 TensorFlow-Probability 自定义损失

Question

我试图在TensorFlow-Probability 中实现一个屏蔽损失函数，它可以忽略标签中的NA。

对于常规张量来说，这是一项非常古老的任务。我找不到发行版的示例。

我的分布大小(批次、时间步长、输出)(512、251 天、1 到 8 个时间序列)

示例中给出的传统损失函数是使用分布的对数概率。

neg_log_likelihood <- function (x, rv_x) {
-1*(rv_x %>% tfd_log_prob(x))
}

当我用零替换 NA 时，模型训练良好并收敛。当我离开 NA 时，它会产生预期的 NaN 损失。

我已经尝试了许多不同的 tf$ 排列，其中用 0 替换损失，用 0 替换标签，等等。在每种情况下，模型都会停止训练，损失保持在某个常数附近。即使标签中只有一个 NA，情况也是如此。

neg_log_likelihood_missing <-  function (x, rv_x) {

loss =     -1*(  rv_x %>% tfd_log_prob(x) ) 

loss_nonan = tf$where( tf$math$is_finite(x) , loss, 0  )

return( 
loss_nonan
)
}

我在这里使用R是偶然的，任何python或其他例子我都可以翻译。如果有正确的方法，使损失正确地反向传播，我将不胜感激。

Answer 1

如果您使用的是基于梯度的推理，则可能需要"双where"技巧。

虽然这会得到正确的值y：

y = computation(x)
tf.where(is_nan(y), 0, y)

。tf.where的导数仍然可以有nan.

而是写：

safe_x = tf.where(is_unsafe(x), some_safe_x, x)
y = computation(safe_x)
tf.where(is_unsafe(x), 0, y)

。既安全y，又安全dy/dx.

对于您正在考虑的情况，也许可以写：

class MyMaskedDist(tfd.Distribution):
...
def _log_prob(self, x):
safe_x = tf.where(tf.is_nan(x), self.mode(), x)
lp = compute_log_prob(safe_x)
lp = tf.where(tf.is_nan(x), tf.zeros([], lp.dtype), lp)
return lp