在 tfp 中训练变分贝叶斯神经网络时，如何分别可视化损失中不同项的演变？

Question

我想使用 tensorflow-probability 来训练一个简单的全连接贝叶斯神经网络。损失由 KL 项和负对数似然项组成。我怎样才能看到它们与 tfp 的单独演变？

我有以下代码：

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

tfk = tf.keras
tfkl = tf.keras.layers
tfpl = tfp.layers
tfd = tfp.distributions

[make some data for a regression task]

input = tfkl.Input(n_features)
x = input
x = tfpl.DenseFlipout(100, activation='relu')(x)
x = tfpl.DenseFlipout(2)(x)
x = tfpl.DistributionLambda(lambda t: tfd.Normal(loc=t[..., :1],
                                                 scale=1e-3 + tf.math.softplus(t[..., 1:])))(x)

model = tfk.Model(input, x)

negloglik = lambda y, rv_y: -rv_y.log_prob(y)

model.compile(optimizer=tf.optimizers.Adam(), loss=negloglik, metrics=['mse'])
history = model.fit(x_train, y_train, epochs=10, validation_data=(x_val, y_val));

损失函数是显式项 negloglik 和每个 DenseFlipout 层中的 KL 散度项的总和（我可以通过查看 model.losses 看到它们在那里，因为例如）。

如何分别可视化这些术语？

一次尝试：

如果我尝试在指标中添加一个计算negloglik的函数，例如

def negloglik_met(y_true, y_pred):
    return -y_pred.log_prob(y_true)

我得到 AttributeError: 'Tensor' object has no attribute 'log_prob'，这让我很困惑。 y_pred应该是DistributionLambda层的输出，那为什么是Tensor而不是Distribution呢？

其他我希望能奏效但没有奏效的方法是将 model.losses[0] 添加到指标中。我得到 ValueError: Could not interpret metric function identifier: Tensor("dense_flipout/divergence_kernel:0", shape=(), dtype=float32).

Answer 1

我深入研究了 TensorFlow 代码。这是由于自动 TensorFlow 在您的 (lambda) 函数周围创建了一个自动包装器。它将模型输出（分布）投射并重塑为度量的类型（无论如何对我来说这似乎很奇怪）。因此，为了防止它，您应该创建自己的包装器，它不执行此转换。执行此操作的代码位于：https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/metrics.py#L583

因此，在该代码块上激发自己的灵感，制作您自己的指标包装器。这应该是TFP的一个特点。

在 tfp 中训练变分贝叶斯神经网络时，如何分别可视化损失中不同项的演变？

When training a variational Bayesian neural network in tfp, how can I visualize the evolution of the different terms in the loss separately?

python

tensorflow

machine-learning

neural-network

tensorflow-probability