Pytorch之KLDivLoss

理论基础

KL散度：衡量两个概率分布之间的相似性，其值越小，概率分布越接近。公式表达如下。

$$\begin{array}{rl}{D}_{KL}(P\Vert Q)& =\sum _{i=1}^N\left[p\left(x_i\right)\log p\left(x_i\right)-p\left(x_i\right)\log q\left(x_i\right)\right]\\ & =\sum _{i=1}^N\left[p\left(x_i\right)\frac{\log p\left(x_i\right)}{\log q\left(x_i\right)}\right]\end{array}$$

注：对于两个概率分布 $P$ 、$Q$，$P$ 为真实事件的概率分布，$Q$ 为随机事件拟合出来的该事件的概率分布，即 $D_{KL}(P\Vert Q)$ 表示使用 $P$ 来拟合 $Q$， 或者说使用 $Q$ 来指导 $P$。

实现

import torch
import torch.nn as nn
import torch.nn.functional as F

# 预测值
input = torch.tensor([0.7, .1, .2], requires_grad=True)  # dim=0 每一行为一个样本

# 真实值
target = torch.tensor([.2, .5, .3])

# 计算KL散度
# 方式1
kl_loss = nn.KLDivLoss(reduction="batchmean")
output = kl_loss(F.log_softmax(input, dim=0), F.softmax(target, dim=0))
print(output)

# 方式2
print(F.kl_div(F.log_softmax(input, dim=0), F.softmax(target, dim=0), reduction="batchmean"))

# 方式3
my_kl_loss = F.softmax(target, dim=0) * (torch.log(F.softmax(target, dim=0)) - F.log_softmax(input, dim=0))
my_kl_loss = my_kl_loss.mean()
print(my_kl_loss)

# 方式4
my_kl_loss2 = F.softmax(target, dim=0) * (F.log_softmax(target, dim=0) - F.log_softmax(input, dim=0))
my_kl_loss2 = my_kl_loss2.mean()
print(my_kl_loss2)

# ----------------输出--------------------
# tensor(0.0239, grad_fn=<DivBackward0>)
# tensor(0.0239, grad_fn=<MeanBackward0>)
# tensor(0.0239, grad_fn=<MeanBackward0>)
# tensor(0.0239, grad_fn=<DivBackward0>)
# ----------------------------------------

几个要点：

1. KL散度的原理
2. KL实现为什么要做log和softmax
3. 上溢出和下溢出的情况
4. 在pytorch的log函数中，默认是以 $e$ 为底数的

---

参考：

• loss函数之KLDivLoss
• KL散度理解以及使用pytorch计算KL散度
• 有效防止softmax计算时上溢出（overflow）和下溢出（underflow）的方法
• 交叉熵