Focal Loss实现

一、序言

        Focal Loss是分类常用到的损失函数，面试中也经常容易考，所以这里记录一下。

二、公式和代码

        公式：

        Focal Loss = -alpha*(1-pt)**gamma*log(pt)

        其中alpha用来调节类别间的平衡，gamma用来平衡难例和简单例的平衡。

        代码：

import torch
import torch.nn as nn
import torch.nn.functional as F


class FocalLoss(nn.Module):
    def __init__(self, alpha=None, gamma=2.0, reduction="mean"):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.reduction = reduction
        if alpha == None:
            self.alpha = torch.tensor([1.0])
        else:
            self.alpha = torch.tensor(alpha)

    def forward(self, predict, label):
        bce_loss = F.cross_entropy(predict, label, reduction="none")
        pt = torch.exp(-bce_loss)
        alpha = self.alpha.gather(0, label.data.view(-1))
        loss = alpha * (1 - pt) ** self.gamma * bce_loss
        return loss.mean()


if __name__ == "__main__":
    focal_loss = FocalLoss(alpha=[0.25, 0.5, 1.0], gamma=2.0)
    predict = torch.randn(10, 3, requires_grad=True)
    label = torch.empty(10, dtype=torch.long).random_(3)
    loss = focal_loss(predict, label)
    print("loss:", loss.item())