torch.nn.CrossEntropyLoss 里面究竟发生了什么

假设

假设有三个分类，模型输出值为 output = model(input)，得到如下输出向量
$$[o_1,o_2,o_3]$$

假设 target 为 2， one hot 编码后为 $[0,0,1]$

使用 torch.nn.CrossEntropyLoss 对 output 和 target 计算损失的过程如下所示：

过程

输出值 output 经过 Softmax 后得到每个类别的预测值

$$[S_1,S_2,S_3]$$

然后预测值再经过 log 后得到 LogSoftmax
$$[\log (S_1),\log (S_2),\log (S_3)]$$

最后再把 LogSoftmax 的结果与 target 的 one hot 值输入 NLLLoss 后得到
$$-\log (S_3)$$
因为 target 的 one hot值 为 [0,0,1]，所以只保留 $\log (S_3)$ 的值，然后取符号

代码验证

import torch
import torch.nn as nn

if __name__ == "__main__":
    
    torch.manual_seed(2020)
    output = torch.randn((1, 3))
    print("output: \t\t", output)

    target = torch.tensor([2])
    print("target: \t\t", target)

    ss = torch.softmax(output, dim=1)
    print("softmax: \t\t", ss)

    logsoftmax = torch.log(ss)
    print("logsoftmax: \t\t", logsoftmax)

    nllloss = -logsoftmax[0][2]
    print("nllloss : \t\t", nllloss)
    
    celoss = nn.CrossEntropyLoss()(output, target)
    print("CrossEntropyLoss : \t", celoss) 

打印输出结果为：

output:                  tensor([[ 1.2372, -0.9604,  1.5415]])
target:                  tensor([2])
softmax:                 tensor([[0.4054, 0.0450, 0.5496]])
logsoftmax:              tensor([[-0.9029, -3.1005, -0.5986]])
nllloss :                tensor(0.5986)
CrossEntropyLoss :       tensor(0.5986)

验证结果正确

总结

根据官方文档的解释：This criterion combines nn.LogSoftmax()and nn.NLLLoss() in one single class. 其中 nn.LogSoftmax 只接收一个参数 output，先执行 softmax 操作，然后再取 log 对数，nn.NLLLoss 则需传入两个参数，一个是 nn.LogSoftmax 的输出 log，一个是 target，将 target one hot 编码后，取 logs 对应 one hot 编码 非 0 位置的值，然后再取负数。