如何使用多元高斯的对数似然在 Keras 中进行自定义损失？

如何解决如何使用多元高斯的对数似然在 Keras 中进行自定义损失？

我正在尝试使用 Keras 复制 conditional neural processes 论文中关于函数回归的第一个实验。

因此，我使用我的观察结果的负对数似然作为损失函数。我的网络输出两个值，每个观察值的平均值和标准差。由于网络可以将可变数量的目标点作为输入，我的批次在解码器的输出端具有形状 [None,None,2]（对于 [batch_size,number_of_target_points,output_size]）。我目前定义自定义损失的方式如下：

def custom_loss(y_true,y_pred):
    mu,sigma = tf.split(y_pred,2,axis=-1)
    dist = tfp.distributions.MultivariatenormalDiag(loc=mu,scale_diag=sigma)
    loss = - tf.reduce_mean(K.log(dist.log_prob(y_true)))
    return loss

opt = Adam(learning_rate=1e-4)
model.compile(loss = custom_loss,optimizer = opt)

但是，我在每个时期的训练过程中都会损失 nan，就像这样

Epoch 1/10
100/100 [==============================] - 7s 57ms/step - loss: nan - val_loss: nan

有人知道我做错了什么以及我该如何解决吗？

请注意，网络末端的方差地板层确保标准偏差为正且不会太小：

def floor_variance(x):
  # Repeat r vector
  mu,log_sigma = tf.split(x,axis=-1)
  sigma = 0.1 + 0.9 * tf.nn.softplus(log_sigma)
  return K.concatenate([mu,sigma],axis=-1)

[... define the network ... ]
x = Lambda(floor_variance,name = 'Floor_variance')(x) # x shape [None,2]