python - TensorFlow 有内置 KL 散度损失函数吗？

Question

我有两个张量，prob_a 和 prob_b，形状为 [None, 1000]，我想从 计算 KL 散度code>prob_a 到 prob_b。 TensorFlow 中是否有针对此的内置函数？我尝试使用 tf.contrib.distributions.kl(prob_a, prob_b)，但它给出了:

NotImplementedError: No KL(dist_a || dist_b) registered for dist_a type Tensor and dist_b type Tensor

如果没有内置函数，什么是好的解决方法？

Answer 1

假设您的输入张量 prob_a 和 prob_b 是沿着最后一个轴总和为 1 的概率张量，您可以这样做:

def kl(x, y):
    X = tf.distributions.Categorical(probs=x)
    Y = tf.distributions.Categorical(probs=y)
    return tf.distributions.kl_divergence(X, Y)

result = kl(prob_a, prob_b)

一个简单的例子:

import numpy as np
import tensorflow as tf
a = np.array([[0.25, 0.1, 0.65], [0.8, 0.15, 0.05]])
b = np.array([[0.7, 0.2, 0.1], [0.15, 0.8, 0.05]])
sess = tf.Session()
print(kl(a, b).eval(session=sess))  # [0.88995184 1.08808468]

你会得到相同的结果

np.sum(a * np.log(a / b), axis=1) 

但是，此实现有点错误(已在 Tensorflow 1.8.0 中检查过)。

如果您在 a 中的概率为零，例如如果您尝试使用 [0.8, 0.2, 0.0] 而不是 [0.8, 0.15, 0.05]，您将得到 nan，即使是 Kullback- Leibler 定义 0 * log(0/b) 应为零。

为了缓解这种情况，应该添加一些小的数值常量。在这种情况下，使用 tf.distributions.kl_divergence(X, Y, allow_nan_stats=False) 会导致运行时错误也是明智的。

此外，如果 b 中有一些零，您将获得 inf 值，allow_nan_stats=False 不会捕获这些值选项，因此也必须处理这些选项。