Focal Loss代码

Focal Loss 是一种用于解决类别不平衡问题的损失函数，其可以使得网络更加关注于难以分类的样本。以下是 Focal Loss 的代码实现：

import torch
import torch.nn.functional as F

class FocalLoss(torch.nn.Module):
    def __init__(self, alpha=1, gamma=2, reduction='mean'):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction

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

其中，`alpha` 表示样本类别的权重系数，`gamma` 是一个调节因子，用于调整难分类样本的权重；`reduction` 参数用于指定损失函数的计算方式，可以选择取平均值或者总和。

使用该损失函数时，只需要将其作为模型的损失函数进行传入即可。例如：

model = MyModel()
criterion = FocalLoss(alpha=1, gamma=2, reduction='mean')
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
for epoch in range(num_epochs):
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()