一文讲清楚【KL距离】、【torch.nn.functional.kl

• 一、定义
• • 1. KL距离(Kullback–Leibler divergence)
  • 2. torch.nn.KLDivLoss()
  • 3. torch.nn.functional.kl_div()
• 二、实验
• • 1. torch.nn.KLDivLoss()
  • 2. torch.nn.functional.kl_div()
• 三、结论
• • 1. 实际应用

设P, Q为两个概率(意思就是sum为1，相当于做完softmax之后的效果)，则KL距离 D K L ( P ∣ ∣ Q ) D_{KL}(P||Q) DKL​(P∣∣Q)被定义为

一个类(class)，通过调用(__call__)来实现loss的计算。官网的定义如下：

这里可以看出来，y相当于Q，x相当于P。想要实现KL距离 D K L ( P ∣ ∣ Q ) D_{KL}(P||Q) DKL​(P∣∣Q)的话，其使用方法为

# P, Q are two probabilities (normally after softmax)
l = torch.nn.KLDivLoss()
loss = l(torch.log(Q), P)

这是很显然的，因为根据官网定义
l n = y n ⋅ ( log ⁡ y n − x n ) = y n ⋅ ( log ⁡ y n − log ⁡ Q n ) = P n ⋅ ( log ⁡ P n − log ⁡ Q n ) = P n ⋅ log ⁡ P n Q n begin{aligned} l_n &= y_n cdot (log y_n - x_n)\ &= y_n cdot (log y_n - log Q_n) \ & = P_ncdot (log P_n - log Q_n) \ & = P_ncdot log frac{P_n}{ Q_n} end{aligned} ln​​=yn​⋅(logyn​−xn​)=yn​⋅(logyn​−logQn​)=Pn​⋅(logPn​−logQn​)=Pn​⋅logQn​Pn​​​

一个函数(function)，通过直接使用来实现loss的计算。想要实现KL距离 D K L ( P ∣ ∣ Q ) D_{KL}(P||Q) DKL​(P∣∣Q)的话，其使用方法为

# P, Q are two probabilities (normally after softmax)
loss = torch.nn.functional.kl_div(torch.log(Q), P)

简单来说，就是torch.nn.KLDivLoss()的简化版本，随调随用。

我们简单设计个小实验

import torch
P = torch.tensor([0.25] * 4 + [0])
Q = torch.tensor([0.2] * 5)
l = torch.nn.KLDivLoss(reduction='sum')
loss = l(torch.log(Q), P)  
print(loss)
# output: tensor(0.2231)

我们用计算器算一下
l o s s = ∑ i = 1 i = 5 P i log ⁡ P i Q i = 4 × 0.25 × ln ⁡ 0.25 0.2 + 0 = 0.22314 loss = sum_{i=1}^{i=5} P_i log frac{P_i}{Q_i} = 4times0.25timeslnfrac{0.25}{0.2} +0= 0.22314 loss=i=1∑i=5​Pi​logQi​Pi​​=4×0.25×ln0.20.25​+0=0.22314
可以看到，是一致的。此外，一般reduction都用batchmean，这个只是最后要不要平均一下罢了，详见官网。

类似的，我们也做一个实验

import torch
P = torch.tensor([0.25] * 4 + [0])
Q = torch.tensor([0.2] * 5)
loss = torch.nn.functional.kl_div(torch.log(Q), P, reduction='sum') 
print(loss) 
# output: tensor(0.2231)

实验成功！

一般地，一个batch的sample通过model后，我们会得到[batch_size, class_num]的矩阵。如果model的最后一层有softmax，我们就相当于拿到了概率Q。

而对应的label(大小为[batch_size, 1])则对应P，但是我们需要先对label进行softmax归一化。所以，伪代码如下：

import torch
import torch.nn.functional as F
...
predicted = model(samples)
log_pre = torch.log(predicted)
labels = F.softmax(labels)
loss = F.kl_div(log_pre, labels, reduction='batchmean') 

类似的，如果model的最后没有softmax，伪代码如下：

import torch
import torch.nn.functional as F
...
predicted = model(samples)
log_pre = nn.LogSoftmax(dim=1)(predicted)
labels = F.softmax(labels)
loss = F.kl_div(log_pre, labels, reduction='batchmean')