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 )