Pytorch学习笔记(8)——在序列标注等多维数据上如何使用交叉熵

最近遇到个新的问题，要对序列标注任务使用交叉熵获得损失，由于没有在网上查找到相关资料，所以就自己整理了一份如何调库的方法。

对于文本分类等任务而言，其模型输出的数据格式为 $(batch\_size,num\_classes)$，这类方法采用 Pytorch 的交叉熵很简单，代码如下：

import torch
import torch.nn as nn


# output shape: torch.Size([4, 2])
output = torch.FloatTensor([[0.8, 0.2],
                            [0.6, 0.4],
                            [0.3, 0.7], 
                            [0.1, 0.9]])
# label shape: torch.Size([4])                          
label = torch.LongTensor([0, 0, 1, 0])

loss = nn.CrossEntropyLoss()
# tensor(0.6799)
loss(output, label)

=================================== 分隔符 ===================================
注：如果同学你选择手算了一遍上面的数据，你会发现你手算的结果和 loss 输出的结果完全不同，无论算多少遍都不对，这是因为如下的情况。

$$\mathcal{L}=-\frac{1}{N}\sum _{i=1}^N\left[y_i\log (p_i)+(1-y_i)\log (1-p_i)\right]$$

如果我们选择手算，将上面的数据代入到上面的公式，即：
$$\mathcal{L}=-\frac{1}{4}(\log (0.2)+\log (0.4)+\log (0.7)+\log (0.9))=0.7469$$
与代码得出的结果不符，这是因为 Pytorch 的 CrossEntropyLoss() 里面自带了 Softmax，Softmax 后的结果如下：

tensor([[0.6457, 0.3543],
        [0.5498, 0.4502],
        [0.4013, 0.5987],
        [0.3100, 0.6900]])

将这个 Softmax 后的结果再代入到公式里面算，算出来就是 0.6799 了。这件事情告诉我们，千万别在做这类任务的时候，在输出层加 softmax 后再拼接 CrossEntropyLoss() 函数，这样会导致损失可能会出现过大的情况，解决方案就是只输出 logits (不接 softmax 的输出层)，再将 logits 丢入到 CrossEntropyLoss() 中。（BTW，我被这个给坑麻了）

=================================== 分隔符 ===================================

好的，回到正文来，对于序列标注等模型输出是高维数据，比如序列标注中是三维数据，即 $(batch\_size,max\_length,num\_classes)$，这类数据如果我们直接丢入到 CrossEntropyLoss() 中会报如下错误：

# output shape: torch.Size([1, 4, 2])
output = torch.FloatTensor([[[0.8, 0.2],
                             [0.6, 0.4],
                             [0.3, 0.7], 
                             [0.1, 0.9]]])
# label shape: torch.Size([1, 4])                          
label = torch.LongTensor([[0, 0, 1, 0]])

loss = nn.CrossEntropyLoss()
loss(output, label)

RuntimeError: Expected target size [1, 2], got [1, 4]

为了解决这个问题，我就查了 torch 的官方文档和源码，发现文档中是这么写的：

而源码中的注释是这样的（为了方便大家阅读，我将其转为了 latex）：

• input (Tensor) : $(N,C)$ where $C=numberofclasses$ or $(N,C,H,W)$ in case of 2D Loss, or $(N,C,d_1,d_2,...,d_K)$ where $K\ge 1$ in the case of K-dimensional loss.
• target (Tensor) : $(N)$ where each value is $0\le \text{targets}[i]\le C-1$, or $(N,d_1,d_2,...,d_K)$ where $K\ge 1$ for K-dimensional loss.

简单来说：

1. 如果你的输出是二维（即 $(batch\_size,num\_classes)$），那么 target 的维度就应该为 $(batch\_size)$；
2. 如果你的输出维度是大于二维的（假设为 $(batch\_size,max\_length,num\_classes)$），那么应该给转置一下，保证第二维为 $num\_classes$。

带着第二个思路，我们将代码给更改为：

# output shape: torch.Size([1, 4, 2])
output = torch.FloatTensor([[[0.8, 0.2],
                             [0.6, 0.4],
                             [0.3, 0.7], 
                             [0.1, 0.9]]])
# label shape: torch.Size([1, 4])                          
label = torch.LongTensor([[0, 0, 1, 0]])

# output shape: torch.Size([1, 2, 4])
output = output.permute(0, 2, 1)
loss = nn.CrossEntropyLoss()

# tensor(0.6799)
loss(output, label)

输出的结果就和二维的输出结果一样了！