python - TensorFlow中的KL分歧

我有两个张量，prob_a和prob_bwith shape[None, 1000]，我想计算从prob_a到prob_b的kl发散。在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
如果没有内置功能，那么什么是一个好的解决方案？

最佳答案

假设输入张量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选项捕获，因此这些值也必须处理。