KLD Loss（ tf.nn.softmax, torch.nn.functional.softmax, log_softmax, kl_div） 计算技巧(一)

最近在比较不同模型的性能，发现虽然文献中使用的相同的指标，比如KLD。但是数据的处理方式却存在着差异，这会导致最后的数据并不具有直接可比性。
这里记录下，其中的一些值得记住的细节。主要涉及的API包括tf.nn.softmax, torch.nn.functional.softmax, log_softmax, kl_div

二维数据输入

为直观的看出数据，我们以一个2x2的矩阵为例，并打印。

import cv2
import numpy as np
import torch
import torch.nn.functional as F
import tensorflow as tf


y_pred=np.random.randn(1,2,2)
print( '\t y_pred', y_pred)

y_pred [[[-1.23909949 0.77767204]
[ 0.08646117 -0.14608897]]]

数据前处理

这里开始就有技巧性，由于神经网络的预测输出数值一般为双极性数值。如何将数据进行合理的处理，使其映射到标准空间方便后续计算。
这里我觉得比较合适的操作方法是先进行一组softmax操作，不管输入数据的范围是多少先映射到[0,1]空间。这里的softmax操作就有技巧，我们先看看softmax的API定义

tf.nn.softmax

tf.nn.softmax(
    logits, axis=None, name=None)

其中的axis形参，默认对最后1维度进行softmax操作

The dimension softmax would be performed on. The default is -1 which indicates the last dimension.

参考：https://www.tensorflow.org/api_docs/python/tf/nn/softmax

因此如果我们直接使用softmax操作，得到的是对最后维度，即 [-1.23909949 0.77767204] 和[ 0.08646117 -0.14608897]分别进行softmax操作的结果

y_pred_soft=tf.nn.softmax(y_pred)
print('tf softmax y_pred:', y_pred_soft)

输出

tf softmax y_pred: tf.Tensor(
[[[0.11745323 0.88254677]
[0.55787694 0.44212306]]], shape=(1, 2, 2), dtype=float64)

torch.nn.functional.softmax

对于pytorch的softmax操作

torch.nn.functional.softmax(input, dim=None, _stacklevel=3, dtype=None)

• input (Tensor) – input
• dim (int) – A dimension along which softmax will be computed.
• dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None
  参考：https://pytorch.org/docs/stable/nn.functional.html

因此对应tf的默认操作，这里pytorch应该写成

y_pred = torch.from_numpy(y_pred)
y_pred = F.softmax(y_pred,dim=-1)
print('torch softmax y_pred:', y_pred)

结果

torch softmax y_pred: tensor([[[0.1175, 0.8825],
[0.5579, 0.4421]]], dtype=torch.float64)

但是但是但是，重点！！！
我们希望的softmax应该是对二维数据中所有元素同时进行的softmax，而不是特定在某个维度，因此这里我们需要的操作，是先将所有数据展开成一维后再进行softmax操作。

y_pred = y_pred.view(1, -1)
y_pred = F.softmax(y_pred, dim=1)

kld 散度计算

tensorflow

按照计算公式：
$$D(p||q)=\sum _{x\in \chi }p(x)log\frac{p(x)}{q(x)}=E_{p}log\frac{p(x)}{q(x)}$$

只需要又以下代码实现即可

kld_loss = tf.reduce_sum(y_true*tf.math.log(y_true/y_pred), axis=1)

pytorch

同理在pytorch中代码实现如下：

kld_loss = (y_true*torch.log(y_true/y_pred)).sum(-1)

但是但是但是，重点！！！
由于pytorch有torch.nn.functional.kl_div，我们可以直接使用。

torch.nn.functional.kl_div(input, target, size_average=None, reduce=None, reduction='mean', log_target=False)

这里又有一个非**常重要的技巧！！！**如果直接将概率y_pred传入input，kld将是一个负值，需要对其取对数，如下进行计算。

kld = F.kl_div(torch.log(y_pred), y_true, reduction='none')

这时候细心的人可能以及注意到了我们对y_input前后进行了softmax和log操作。其实在torch中有F.log_softmax函数，可以直接进行计算。
这里关于log_softmax想多说一点。因为不太理解为什么要使用log_softmax，在网上查了一查觉得这两位博主写得很有道理。

• log_softmax能够解决函数overflow和underflow，加快运算速度，提高数据稳定性。
• 方便两类的分类问题，即预测输出为一维（0或1）的CrossEntropyLoss的计算

CrossEntropyLoss的批注This criterion combines nn.LogSoftmax() and nn.NLLLoss() in one single class)

https://www.zhihu.com/question/358069078
https://zhuanlan.zhihu.com/p/95415762
pytorch中 log_softmax源代码的注解

def log_softmax(input, dim=None, _stacklevel=3, dtype=None):
    # type: (Tensor, Optional[int], int, Optional[int]) -> Tensor
    r"""Applies a softmax followed by a logarithm.

    While mathematically equivalent to log(softmax(x)), doing these two
    operations separately is slower, and numerically unstable. This function
    uses an alternative formulation to compute the output and gradient correctly.

    See :class:`~torch.nn.LogSoftmax` for more details.

    Arguments:
        input (Tensor): input
        dim (int): A dimension along which log_softmax will be computed.
        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.
          If specified, the input tensor is casted to :attr:`dtype` before the operation
          is performed. This is useful for preventing data type overflows. Default: None.
    """

因此对于y_pred的数值，直接使用log_softmax可以大大提高计算效率，并避免由于q(x)为0出现的数值计算问题。
$$D(p||q)=\sum _{x\in \chi }p(x)log\frac{p(x)}{q(x)}=E_{p}log\frac{p(x)}{q(x)}$$

完整代码

import cv2
import numpy as np
import torch
import torch.nn.functional as F
import tensorflow as tf

y_true=np.random.randn(1,2,2)
y_pred=np.random.randn(1,2,2)
print('y_true',y_true, '\t y_pred', y_pred)

def tf_kld (y_true, y_pred):
    
    y_true = tf.convert_to_tensor(y_true)
    y_pred = tf.convert_to_tensor(y_pred)
    y_true = tf.reshape(y_true, [1,-1])
    y_pred = tf.reshape(y_pred, [1,-1])
    y_true=tf.nn.softmax(y_true)
    y_pred=tf.nn.softmax(y_pred)
    print('tf softmax y_pred:', y_pred)
    kld_loss = tf.reduce_sum(y_true*tf.math.log(y_true/y_pred), axis=1)
    return kld_loss

def torch_kld(y_true, y_pred):
    y_true = torch.from_numpy(y_true)
    y_pred = torch.from_numpy(y_pred)
    y_true = y_true.view(1, -1)
    y_pred = y_pred.view(1, -1)
    y_true = F.softmax(y_true, dim=1)
    y_pred = F.log_softmax(y_pred, dim=1)
    print('torch softmax y_pred:', y_pred)
    # print('torch kld loss:', (y_true*torch.log(y_true/y_pred)).sum(-1))
    kld = F.kl_div(y_pred, y_true, reduction='none')
    kld_loss = kld.sum(-1)
    return kld_loss

print('tensorflow kld loss:', tf_kld(y_true, y_pred))
print('pytorch kld loss:', torch_kld(y_true, y_pred))