[pytorch] 图像识别之focal loss (+ohem)

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
class FocalLoss(nn.Module):
   def __init__(self, class_num, alpha=None, gamma=2, size_average=True):
       super(FocalLoss, self).__init__()
       if alpha is None:
           self.alpha = Variable(torch.ones(class_num, 1))
       else:
           if isinstance(alpha, Variable):
               self.alpha = alpha
           else:
               self.alpha = Variable(alpha)
       self.gamma = gamma
       self.class_num = class_num
       self.size_average = size_average
   def forward(self, inputs, targets):
       N = inputs.size(0)
       C = inputs.size(1)
       P = F.softmax(inputs)