PyTorch中两个概率分布的KL发散

Question

我有两个概率分布。如何在PyTorch中找到它们之间的KL差异？常规互熵仅接受整数标签。

Answer 1

是的，PyTorch有一个名为方法kl_div下torch.nn.functional到张量之间的直接计算KL-devergence。假设您有张量a并且b形状相同。您可以使用以下代码：

import torch.nn.functional as F
out = F.kl_div(a, b)

有关更多详细信息，请参见上面的方法文档。

Answer 2

functionkl_div与wiki的解释不同。

我使用以下内容：

# this is the same example in wiki
P = torch.Tensor([0.36, 0.48, 0.16])
Q = torch.Tensor([0.333, 0.333, 0.333])

(P * (P / Q).log()).sum()
# tensor(0.0863), 10.2 µs ± 508

F.kl_div(Q.log(), P, None, None, 'sum')
# tensor(0.0863), 14.1 µs ± 408 ns

相比kl_div，甚至更快

Answer 3

如果你有两个pytorch distributionobject形式的概率分布。那么您最好使用该功能torch.distributions.kl.kl_divergence(p, q)。有关文档，请点击链接