【深度学习】pytorch计算KL散度、kl_div

忘却的旋律dw2024-04-13 12:58

使用pytorch进行KL散度计算，可以使用pytorch的kl_div函数

假设y为真实分布，x为预测分布。

java 复制代码

import torch
import torch.nn.functional as F

# 定义两组数据
tensor1 = torch.tensor([[0.1, 0.2, 0.3, 0.2, 0.2],
                        [0.2, 0.1, 0.2, 0.3, 0.2],
                        [0.2, 0.3, 0.1, 0.2, 0.2],
                        [0.2, 0.2, 0.3, 0.1, 0.2],
                        [0.2, 0.2, 0.2, 0.3, 0.1],
                        [0.1, 0.2, 0.2, 0.2, 0.3],
                        [0.3, 0.2, 0.1, 0.2, 0.2],
                        [0.2, 0.3, 0.2, 0.1, 0.2],
                        [0.1, 0.2, 0.2, 0.3, 0.2],
                        [0.2, 0.1, 0.3, 0.2, 0.2],
                        [0.2, 0.3, 0.2, 0.2, 0.1],
                        [0.1, 0.1, 0.2, 0.3, 0.3],
                        [0.3, 0.2, 0.2, 0.1, 0.2],
                        [0.2, 0.3, 0.1, 0.2, 0.2],
                        [0.1, 0.3, 0.2, 0.2, 0.2],
                        [0.2, 0.2, 0.1, 0.3, 0.2]])

tensor2 = torch.tensor([[0.2, 0.1, 0.3, 0.2, 0.2],
                        [0.3, 0.2, 0.2, 0.1, 0.2],
                        [0.2, 0.3, 0.2, 0.2, 0.1],
                        [0.1, 0.2, 0.3, 0.2, 0.2],
                        [0.2, 0.2, 0.1, 0.2, 0.3],
                        [0.3, 0.2, 0.2, 0.3, 0.0],
                        [0.2, 0.3, 0.1, 0.2, 0.2],
                        [0.1, 0.2, 0.2, 0.3, 0.2],
                        [0.2, 0.1, 0.3, 0.2, 0.2],
                        [0.2, 0.3, 0.2, 0.1, 0.2],
                        [0.1, 0.2, 0.3, 0.2, 0.2],
                        [0.2, 0.3, 0.2, 0.2, 0.1],
                        [0.2, 0.1, 0.2, 0.3, 0.2],
                        [0.3, 0.2, 0.2, 0.1, 0.2],
                        [0.2, 0.2, 0.3, 0.2, 0.1],
                        [0.1, 0.3, 0.2, 0.2, 0.2]])

# 计算两组张量之间的 KL 散度
logp_x = F.log_softmax(tensor1, dim=-1)
p_y = F.softmax(tensor2, dim=-1)

kl_divergence = F.kl_div(logp_x, p_y, reduction='batchmean')
kl_sum = F.kl_div(logp_x, p_y, reduction='sum')
print("KL散度(batchmean)值为:", kl_divergence.item())
print("KL散度(sum)值为:", kl_sum.item())

打印结果：

复制代码

KL散度(batchmean)值为: 0.00508523266762495
KL散度(sum)值为: 0.0813637226819992

其中kl_div接收三个参数，第一个为预测分布，第二个为真实分布，第三个为reduction。(其实还有其他参数，只是基本用不到)

这里有一些细节需要注意，第一个参数与第二个参数都要进行softmax(dim=-1)，目的是使两个概率分布的所有值之和都为1，若不进行此操作，如果x或y概率分布所有值的和大于1，则可能会使计算的KL为负数。

softmax接收一个参数dim，dim=-1表示在最后一维进行softmax操作。

除此之外，第一个参数还要进行log()操作（至于为什么，大概是为了方便pytorch的代码组织，pytorch定义的损失函数都调用handle_torch_function函数，方便权重控制等），才能得到正确结果。还有说是因为要用y指导x，所以求x的对数概率，y的概率