Poly1CrossEntropyLoss的pytorch实现

代码改自github

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class Poly1CrossEntropyLoss(nn.Module):
    def __init__(self,
                 num_classes: int,
                 epsilon: float = 1.0,
                 reduction: str = "none"):
        """
        Create instance of Poly1CrossEntropyLoss
        :param num_classes:
        :param epsilon:
        :param reduction: one of none|sum|mean, apply reduction to final loss tensor
        """
        super(Poly1CrossEntropyLoss, self).__init__()
        self.num_classes = num_classes
        self.epsilon = epsilon
        self.reduction = reduction
        return

    def forward(self, output, labels):
        """
        Forward pass
        :param output: tensor of shape [N, num_classes]
        :param labels: tensor of shape [N] or [N, num_classes]
        :return: poly cross-entropy loss
        """
        #timm如果有使用cutmix或者mixup，那么就会得到的labels是[batchsize，classnums]维度的矩阵，因为获取的label不是hard label,但也不算是smooth label，是每个数据对应的label类似于这种[0,0,0.65,0,0,0.35,0,...]
        if labels.ndim == 1:
        	labels_onehot = F.one_hot(labels, num_classes=self.num_classes).to(device=output.device,dtype=output.dtype)
        else:
        	labels=labels.to(device=output.device)
        pt = torch.sum(labels * F.softmax(output, dim=-1), dim=-1)#64,1
        CE = F.cross_entropy(input=output, target=labels, reduction='none')
        # print(CE.shape)#64
        poly1 = CE + self.epsilon * (1 - pt)
        if self.reduction == "mean":
            poly1 = poly1.mean()
        elif self.reduction == "sum":
            poly1 = poly1.sum()
        #加上这个平均操作是为了获取每个数据平均的loss，如果不加就是长度为batchsize的一维矩阵，分别存放每个数据对应的loss
        #poly1 = poly1.mean()
        return poly1