如何在 Pytorch 中计算两批分布之间的 KL 散度？答案

【问题标题】：How to calculate KL Divergence between two batches of distributions in Pytroch?如何在 Pytorch 中计算两批分布之间的 KL 散度？
【发布时间】：2021-04-12 05:09:00
【问题描述】：

给定一批分布，表示为 pytorch 张量：

A = torch.tensor([[0., 0., 0., 0., 1., 5., 1., 2.],
        [0., 0., 1., 0., 4., 2., 1., 1.],
        [0., 0., 1., 1., 0., 5., 1., 1.],
        [0., 1., 1., 0., 2., 3., 1., 1.],
        [0., 0., 2., 1., 3., 1., 1., 0.],
        [0., 0., 2., 0., 5., 0., 1., 0.],
        [0., 2., 1., 4., 0., 0., 1., 0.],
        [0., 0., 2., 4., 1., 0., 1., 0.]], device='cuda:0')

A 是一批分布，由八个分布组成。现在，给定另一批分布B：

B = torch.tensor([[0., 0., 1., 4., 2., 1., 1., 0.],
        [0., 0., 0., 5., 1., 2., 1., 0.],
        [0., 0., 0., 4., 2., 3., 0., 0.],
        [0., 0., 1., 7., 0., 0., 1., 0.],
        [0., 0., 1., 2., 4., 0., 1., 1.],
        [0., 0., 1., 3., 1., 3., 0., 0.],
        [0., 0., 1., 4., 1., 0., 2., 0.],
        [1., 0., 1., 5., 0., 1., 0., 0.],
        [0., 1., 5., 1., 0., 0., 1., 0.],
        [0., 0., 3., 2., 2., 0., 1., 0.],
        [0., 2., 4., 0., 1., 0., 1., 0.],
        [1., 0., 4., 1., 1., 1., 0., 0.]], device='cuda:0')

B 有 12 个分布。我想计算A中的每个分布与B中的每个分布之间的KL散度，然后得到一个KL距离矩阵，其形状为12*8。我知道使用循环结构和torch.nn.functional.kl_div() 来达到它。 pytorch 中是否还有其他方法可以不使用 for-loop 来实现它？

这是我使用 for 循环的实现：

p_1 = F.softmax(A, dim = -1)
p_2 = F.softmax(B, dim = -1)
C = torch.empty(size = (A.shape[0], B.shape[0]), dtype = torch.float)

for i,a in enumerate(p_1):
    for j,b in enumerate(p_2):
        C[i][j] =  torch.nn.functional.kl_div(a.log(), b)
print(C)

输出是：

tensor([[0.4704, 0.5431, 0.3422, 0.6284, 0.3985, 0.2003, 0.4925, 0.5739, 0.5793,
         0.3992, 0.5007, 0.4934],
        [0.3416, 0.4518, 0.2950, 0.5263, 0.0218, 0.2254, 0.3786, 0.4747, 0.3626,
         0.1823, 0.2960, 0.2937],
        [0.3845, 0.4306, 0.2722, 0.5022, 0.4769, 0.1500, 0.3964, 0.4556, 0.4609,
         0.3396, 0.4076, 0.3933],
        [0.2862, 0.3752, 0.2116, 0.4520, 0.1307, 0.1116, 0.3102, 0.3990, 0.2869,
         0.1464, 0.2164, 0.2225],
        [0.1829, 0.2674, 0.1763, 0.3227, 0.0244, 0.1481, 0.2067, 0.2809, 0.1675,
         0.0482, 0.1271, 0.1210],
        [0.4359, 0.5615, 0.4427, 0.6268, 0.0325, 0.4160, 0.4749, 0.5774, 0.3492,
         0.2093, 0.3015, 0.3014],
        [0.0235, 0.0184, 0.0772, 0.0286, 0.3462, 0.1461, 0.0142, 0.0162, 0.3524,
         0.1824, 0.2844, 0.2988],
        [0.0097, 0.0171, 0.0680, 0.0284, 0.2517, 0.1374, 0.0082, 0.0148, 0.2403,
         0.1058, 0.2100, 0.1978]], device='cuda:0')

【问题讨论】：

A 和 B 中的值是“频率”还是“对数概率”？你的概率为零，我想这对 KL 来说可能是个问题。您能否提供您当前使用循环计算的所需输出？
@Shai 感谢您的提醒，我已经修改了问题。

标签： python machine-learning pytorch

【解决方案1】：

看nn.KLDivLoss，计算KL散度的公式是

kl = torch.mean(b * (torch.log(b) - a))

我们可以使用broadcasting 来高效地计算 KL：

# avoid NaNs from log(0)
lB = B.clone()
lB[B==0] = 1.

# do the computation efficiently
C = (B[None, ...] * (torch.log(lB[None, ...]) - A[:, None, :])).mean(dim=-1)

想来想去，我不确定你问的是否有道理。您的 A 和 B 张量充满了数字，但它们不代表分布（它们的总和不等于 1）。请仔细考虑您要在这里做什么。

【讨论】：

感谢您的回答。在我的问题中，A 中的[1,1,2,2,3,5] 表示 1 出现两次，2 出现两次，3 出现一次，5 出现一次。例如A包含三个图（每行代表一个图），每个图中有6个节点。 [1,1,2,2,3,5] 表示A 中第一个图中每个节点的度数。所以我把它当作a分布，我想计算A中的每个度数分布和B中的每个分布之间的KL散度。也许我需要先将它们转换为概率分布。
@MarioKZZ pytorch 解释 A 和 B 的方式与您描述的完全不同。因此，您得到的数字毫无意义。