PyTorch 中的交叉熵函数 CrossEntropyLoss 的计算过程

CrossEntropyLoss() 函数联合调用了 nn.LogSoftmax() 和 nn.NLLLoss()。

假设网络得到的输出为 $h$，它的维度大小为 $B\times C$，其中 $B$ 是 batch_size，$C$ 是分类的总数目。与之对应的训练数据的标签 $y$ 维度是 $1\times B$，$y$ 中元素的取值范围是 $[0,C-1]$，即
$$0\le y[j]\le C-1j=0,1,\cdots ,B-1$$

我们将CrossEntropyLoss() 函数的计算过程拆解为如下两个步骤：

1. 对输出 $h$，执行LogSoftmax(dim=1)，得到 $s$，维度仍然是 $B\times C$。
2. 对 $s$ 执行 $-\log ()$操作，得到负对数概率 $p$，维度仍然是 $B\times C$。

则交叉熵的计算公式为：
(1) L = 1 B ∑ i = 0 B { − log ⁡ ( p [ i , y [ i ] ] ) } L = \frac{1}{B}\sum_{i=0}^B\left\{-\log(p[i,y[i]])\right\} \tag{1} L=B1​i=0∑B​{ −log(p[i,y[i]])}(1)

式(1)其实是从式(2)化简得来的：
$$L=\frac{1}{B}\sum _{i=0}^B\left\{-\sum _{j=0}^{C-1}y[i,j]\log (p[i,j])\right\}$$