import torch
import torch.nn as nn
import torch.nn.functional as F
class CELossWithLabelSmoothing(nn.Module):
''' Cross Entropy Loss with Label Smoothing '''
def __init__(self, label_smoothing=0.1):
super().__init__()
self.label_smoothing = label_smoothing
def forward(self, pred, target):
num_classes = pred.size(1)
print(target)
# 使用标签平滑
one_hot = torch.full((pred.size(0), num_classes), self.label_smoothing / (num_classes - 1))
one_hot.scatter_(1, target.unsqueeze(-1), 1 - self.label_smoothing)
print(one_hot)
# 计算交叉熵损失
loss = -torch.sum(one_hot * torch.log_softmax(pred, dim=1), dim=1)
return loss.mean()
创建一个示例模型输出和真实标签
pred = torch.tensor([[1.0, 2.0, 0.5], [0.5, 1.5, 2.0], [2.0, 1.0, 0.5]])
target = torch.tensor([0, 1, 2])
使用自定义损失函数
criterion = CELossWithLabelSmoothing(label_smoothing=0.1)
loss = criterion(pred, target)
打印损失
print("损失:", loss.item())