把criterion = nn.CrossEntropyLoss().to(device)改成focal loss函数

可以通过以下代码将 criterion 改为 Focal Loss 函数：

import torch.nn.functional as F

class FocalLoss(nn.Module):
    def __init__(self, gamma=2.0, alpha=None, reduction='mean'):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = reduction

    def forward(self, input, target):
        ce_loss = F.cross_entropy(input, target, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = ((1 - pt) ** self.gamma * ce_loss).mean()
        if self.alpha is not None:
            alpha = torch.tensor(self.alpha, device=input.device)
            focal_loss *= F.multi_label_soft_margin_loss(alpha * input.sigmoid(), target, reduction='mean')
        if self.reduction == 'mean':
            return focal_loss
        elif self.reduction == 'sum':
            return focal_loss.sum()
        else:
            return focal_loss

criterion = FocalLoss().to(device)

请注意，这里我们定义了一个新的 FocalLoss 类，并在此类中实现了 Focal Loss 函数。在这个函数中，我们首先计算交叉熵（CE）损失，然后计算 PT，最后使用 Focal Loss 公式计算损失。此外，我们还支持使用 alpha 权重以及不同的 reduction 选项。

最后，我们将 criterion 设置为 Focal Loss 并传入我们定义的设备（device）。