PyTorch交叉熵损失源码与log_softmax实现

脚本专家

PyTorch的nn.CrossEntropyLoss常被直接用于多分类损失计算，但从源码看，它并没有单独实现交叉熵，而是把log_softmax和nll_loss串起来。理解这条调用链，可以解释为什么F.cross_entropy、F.nll_loss(F.log_softmax(...))和nn.CrossEntropyLoss在相同输入上会得到相同loss，也能帮助核对输入维度、target索引和reduction参数。

一、源码调用链
CrossEntropyLoss继承_WeightedLoss，构造参数包括weight、size_average、ignore_index=-100、reduce和reduction='mean'。forward阶段把weight、ignore_index、reduction传给F.cross_entropy：

2. class CrossEntropyLoss(_WeightedLoss):
3.     def __init__(self, weight=None, size_average=None, ignore_index=-100,
4.                  reduce=None, reduction='mean'):
5.         super(CrossEntropyLoss, self).__init__(weight, size_average, reduce, reduction)
6.         self.ignore_index = ignore_index
8.     @weak_script_method
9.     def forward(self, input, target):
10.         return F.cross_entropy(input, target, weight=self.weight,
11.                                ignore_index=self.ignore_index, reduction=self.reduction)

复制代码

F.cross_entropy内部会处理旧版size_average/reduce参数，再执行log_softmax与nll_loss的组合：

2. def cross_entropy(input, target, weight=None, size_average=None, ignore_index=-100,
3.                   reduce=None, reduction='mean'):
4.     if size_average is not None or reduce is not None:
5.         reduction = _Reduction.legacy_get_string(size_average, reduce)
6.     return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)

复制代码

关键点是log_softmax(input, 1)和nll_loss(...)组合，dim=1表示在类别维度归一化。因此CrossEntropyLoss的输入为网络原始输出logits，target为类别索引。从调用链看，传入的input会再次经过log_softmax，排查时可用F.log_softmax+F.nll_loss与F.cross_entropy做等价核对。参数方面，weight用于类别权重，ignore_index默认-100，reduction默认mean，size_average和reduce是旧版兼容入口。

二、为什么需要softmax
交叉熵描述两个概率分布的差异。神经网络输出是向量，不是概率分布；softmax把向量归一化成概率分布，再计算交叉熵。CrossEntropyLoss内部采用log_softmax，即先softmax再取log，然后交给nll_loss。

三、验证三种调用方式等价

2. import torch
3. import torch.nn.functional as F
4. import torch.nn as nn
6. torch.manual_seed(2019)
7. output = torch.randn(2, 3) # 网络输出
8. target = torch.ones(2, dtype=torch.long).random_(3) # 真实标签
9. print('output:\n{}'.format(output))
10. print('target:\n{}\n'.format(target))
12. log_softmax = F.log_softmax(output, dim=1)
13. print('log_softmax:\n{}\n'.format(log_softmax))
15. nll_loss = F.nll_loss(log_softmax, target)
16. print('nll_loss:{}'.format(nll_loss))
18. cross_entropy_loss = F.cross_entropy(output, target)
19. print('cross_entropy_loss:{}'.format(cross_entropy_loss))
21. criterion = nn.CrossEntropyLoss()
22. loss = criterion(output, target)
23. print('loss:{}'.format(loss))

复制代码

输出结果如下：

2. output:
3. tensor([[-0.1187, 0.2110, 0.7463],
4.         [-0.6136, -0.1186, 1.5565]])
5. target:
6. tensor([2, 0])
7. log_softmax:
8. tensor([[-1.5614, -1.2317, -0.6964],
9.         [-2.4335, -1.9386, -0.2635]])
10. nll_loss:1.564985990524292
11. cross_entropy_loss:1.564985990524292
12. loss:1.564985990524292

复制代码

F.log_softmax+ F.nll_loss、F.cross_entropy以及nn.CrossEntropyLoss得到相同值1.564985990524292，说明它们走的是同一计算路径。

四、log_softmax的展开与复现
log_softmax先做softmax，再对结果取log。下面的代码分别用PyTorch和NumPy复现：

2. import torch
3. import torch.nn.functional as F
4. import numpy as np
6. torch.manual_seed(2019)
7. output = torch.randn(2, 3) # 网络输出
8. print(output, '\n')
10. torch_softmax = F.softmax(output, dim=1)
11. print(torch_softmax)
13. torch_log_softmax = torch.log(torch_softmax)
14. print(torch_log_softmax)
16. torch_log_softmax = F.log_softmax(output, dim=1)
17. print(torch_log_softmax)
19. torch_log_softmax = output.log_softmax(dim=1)
20. print(torch_log_softmax, '\n')
22. e_j = np.exp(output.numpy())
23. e_k = e_j.sum(axis=1).reshape(-1, 1)
24. np_softmax = e_j / e_k
25. print(np_softmax)
27. np_log_softmax = np.log(np_softmax)
28. print(np_log_softmax)

复制代码

输出结果如下：

2. tensor([[-0.1187, 0.2110, 0.7463],
3.         [-0.6136, -0.1186, 1.5565]])
4. tensor([[0.2098, 0.2918, 0.4984],
5.         [0.0877, 0.1439, 0.7684]])
6. tensor([[-1.5614, -1.2317, -0.6964],
7.         [-2.4335, -1.9386, -0.2635]])
8. tensor([[-1.5614, -1.2317, -0.6964],
9.         [-2.4335, -1.9386, -0.2635]])
10. tensor([[-1.5614, -1.2317, -0.6964],
11.         [-2.4335, -1.9386, -0.2635]])
12. [[0.20983616 0.29179907 0.49836478]
13. [0.08772492 0.14391233 0.76836276]]
14. [[-1.5614283 -1.2316898 -0.696423 ]
15. [-2.4335492 -1.938551 -0.2634933]]

复制代码

torch.log(F.softmax(output, dim=1))、F.log_softmax(output, dim=1)和output.log_softmax(dim=1)结果一致；NumPy按行计算exp、求和、除法再取log，也能复现相同概率分布和对数结果。

五、使用时的检查点
结合源码参数，使用CrossEntropyLoss时应关注：input形状为[N, C]并在dim=1上做log_softmax；target使用torch.long类别索引，示例中为torch.ones(2, dtype=torch.long).random_(3)；ignore_index默认-100用于指定忽略目标；reduction='mean'决定默认返回均值；weight可传类别权重。排查loss不一致时，先比较F.log_softmax+F.nll_loss与F.cross_entropy，确认调用路径和dim设置，再检查target索引范围与ignore_index、reduction参数是否符合预期。