Pytorch 實現focal_loss 多類別和二分類示例
阿新 • • 發佈:2020-01-15
我就廢話不多說了,直接上程式碼吧!
import numpy as np import torch import torch.nn as nn import torch.nn.functional as F # 支援多分類和二分類 class FocalLoss(nn.Module): """ This is a implementation of Focal Loss with smooth label cross entropy supported which is proposed in 'Focal Loss for Dense Object Detection. (https://arxiv.org/abs/1708.02002)' Focal_Loss= -1*alpha*(1-pt)^gamma*log(pt) :param num_class: :param alpha: (tensor) 3D or 4D the scalar factor for this criterion :param gamma: (float,double) gamma > 0 reduces the relative loss for well-classified examples (p>0.5) putting more focus on hard misclassified example :param smooth: (float,double) smooth value when cross entropy :param balance_index: (int) balance class index,should be specific when alpha is float :param size_average: (bool,optional) By default,the losses are averaged over each loss element in the batch. """ def __init__(self,num_class,alpha=None,gamma=2,balance_index=-1,smooth=None,size_average=True): super(FocalLoss,self).__init__() self.num_class = num_class self.alpha = alpha self.gamma = gamma self.smooth = smooth self.size_average = size_average if self.alpha is None: self.alpha = torch.ones(self.num_class,1) elif isinstance(self.alpha,(list,np.ndarray)): assert len(self.alpha) == self.num_class self.alpha = torch.FloatTensor(alpha).view(self.num_class,1) self.alpha = self.alpha / self.alpha.sum() elif isinstance(self.alpha,float): alpha = torch.ones(self.num_class,1) alpha = alpha * (1 - self.alpha) alpha[balance_index] = self.alpha self.alpha = alpha else: raise TypeError('Not support alpha type') if self.smooth is not None: if self.smooth < 0 or self.smooth > 1.0: raise ValueError('smooth value should be in [0,1]') def forward(self,input,target): logit = F.softmax(input,dim=1) if logit.dim() > 2: # N,C,d1,d2 -> N,m (m=d1*d2*...) logit = logit.view(logit.size(0),logit.size(1),-1) logit = logit.permute(0,2,1).contiguous() logit = logit.view(-1,logit.size(-1)) target = target.view(-1,1) # N = input.size(0) # alpha = torch.ones(N,self.num_class) # alpha = alpha * (1 - self.alpha) # alpha = alpha.scatter_(1,target.long(),self.alpha) epsilon = 1e-10 alpha = self.alpha if alpha.device != input.device: alpha = alpha.to(input.device) idx = target.cpu().long() one_hot_key = torch.FloatTensor(target.size(0),self.num_class).zero_() one_hot_key = one_hot_key.scatter_(1,idx,1) if one_hot_key.device != logit.device: one_hot_key = one_hot_key.to(logit.device) if self.smooth: one_hot_key = torch.clamp( one_hot_key,self.smooth,1.0 - self.smooth) pt = (one_hot_key * logit).sum(1) + epsilon logpt = pt.log() gamma = self.gamma alpha = alpha[idx] loss = -1 * alpha * torch.pow((1 - pt),gamma) * logpt if self.size_average: loss = loss.mean() else: loss = loss.sum() return loss class BCEFocalLoss(torch.nn.Module): """ 二分類的Focalloss alpha 固定 """ def __init__(self,alpha=0.25,reduction='elementwise_mean'): super().__init__() self.gamma = gamma self.alpha = alpha self.reduction = reduction def forward(self,_input,target): pt = torch.sigmoid(_input) alpha = self.alpha loss = - alpha * (1 - pt) ** self.gamma * target * torch.log(pt) - \ (1 - alpha) * pt ** self.gamma * (1 - target) * torch.log(1 - pt) if self.reduction == 'elementwise_mean': loss = torch.mean(loss) elif self.reduction == 'sum': loss = torch.sum(loss) return loss
以上這篇Pytorch 實現focal_loss 多類別和二分類示例就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支援我們。