代码改自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