UnbiasedLossFunction

from __future__ import print_function

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable


class CrossEntropy(nn.Module):
    def __init__(self, alpha=.5):
        super(CrossEntropy, self).__init__()
        self.alpha = alpha

    def forward(self, y_pred, y_actual, size_average=True):
        # sample_weight = y_actual * self.alpha * (1 - y_pred).pow(self.gamma) + (1 - y_actual) * (1 - self.alpha) * y_pred.pow(self.gamma)
        # sample_weight = Variable(sample_weight.data)
        criterion = nn.BCELoss(weight=sample_weight, size_average=size_average)
        loss = criterion(y_pred, y_actual)
        return loss


class FocalLoss(nn.Module):
    def __init__(self, gamma, alpha=.5, reduce=False):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.alpha = alpha
        self.reduce = reduce

    def forward(self, y_pred, y_actual, size_average=True):
        sample_weight = y_actual * self.alpha * (1 - y_pred).pow(self.gamma) + (1 - y_actual) * (1 - self.alpha) * y_pred.pow(self.gamma)
        sample_weight = Variable(sample_weight.data)
        criterion = nn.BCELoss(weight=sample_weight, size_average=size_average, reduce=self.reduce)
        loss = criterion(y_pred, y_actual)
        return loss


class GHMLoss(nn.Module):
    def __init__(self, bins=10):
        super(GHMLoss, self).__init__()
        self.bins = bins
        self.intervals = [float(x) / bins for x in range(bins+1)]
        self.intervals[-1] += 1e-6
        print('GHM intervals\t', self.intervals)

    def forward(self, y_pred, y_actual, size_average=True):
        margin = torch.abs(y_actual - y_pred)#.data.numpy()
        sample_weight = torch.zeros_like(y_actual)
        total = torch.FloatTensor([y_actual.size()[0]])
        for left,right in zip(self.intervals[:-1], self.intervals[1:]):
            index = (margin >= left) & (margin < right) & 1
            if sum(index).float() > 0:
                beta = total / sum(index).float()
                sample_weight[index] = beta
        criterion = nn.BCELoss(weight=sample_weight, size_average=size_average)
        loss = criterion(y_pred, y_actual)
        return loss