分割nn.CrossEntropyLoss()损失如何用mask代码

在使用交叉熵损失函数时，我们可以使用mask来处理填充的位置。对于一个batch中的每个序列，我们可以使用一个长度为该序列长度的mask来标识该序列中哪些位置需要计算损失。具体地，对于填充位置，我们可以将其对应的mask值设为0，而对于非填充位置，我们可以将其对应的mask值设为1。这样，在计算损失时，我们只需要将非填充位置的损失相加再除以非填充位置的数量即可得到平均损失。

以下是使用PyTorch实现分割CrossEntropyLoss()损失时使用mask的示例代码：

import torch
import torch.nn as nn

# 假设我们有一个大小为(batch_size, seq_len, num_classes)的输出张量logits和一个大小为(batch_size, seq_len)的目标张量target
logits = torch.randn(2, 5, 3)
target = torch.tensor([[1, 2, 0, 0, 0], [2, 1, 0, 0, 0]])

# 计算mask
mask = (target > 0).float()

# 定义交叉熵损失函数并计算损失
criterion = nn.CrossEntropyLoss(reduction='none')
loss = criterion(logits.view(-1, logits.size(-1)), target.view(-1))
loss = loss.view(logits.size()[:-1]) * mask
avg_loss = torch.sum(loss) / torch.sum(mask)

print(avg_loss)

其中，`mask`是通过判断`target`是否大于0得到的。在计算损失时，我们先将`loss`按序列中每个位置展开成一维张量，然后再将其reshape成与`logits`相同的形状，并乘以`mask`来过滤掉填充位置。最后，我们将非填充位置的损失相加并除以非填充位置的数量得到平均损失。