loss.py

import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function, Variable

def cross_entropy2d(input, target, weight=None, size_average=True):
    
    n, c, h, w = input.size()
    nt, ct, ht, wt = target.size()
    '''
    # Handle inconsistent size between input and target
    if h > ht and w > wt: # upsample labels
        target = target.unsequeeze(1)
        target = F.upsample(target, size=(h, w), mode='nearest')
        target = target.sequeeze(1)
    elif h < ht and w < wt: # upsample images
        input = F.upsample(input, size=(ht, wt), mode='bilinear')
    elif h != ht and w != wt:
        raise Exception("Only support upsampling")
    '''
    log_p = F.log_softmax(input, dim=1)
    log_p = log_p.transpose(1, 2).transpose(2, 3).contiguous().view(-1, c)
    log_p = log_p[target.contiguous().view(-1, 1).repeat(1, c) >= 0]
    log_p = log_p.view(-1, c)

    mask = target >= 0
    target = target[mask]
    loss = F.nll_loss(log_p, target, ignore_index=250,
                      weight=weight, size_average=False)
    if size_average:
        loss /= mask.data.sum().float()
    return loss

def loss_ce_t(input,target):
    #input=F.sigmoid(input)
    target_bin=Variable(torch.zeros(1,11,target.shape[2],target.shape[3]).cuda().scatter_(1,target,1))
    return F.binary_cross_entropy_with_logits(input,target_bin)

def dice_loss(input, target):
    target_bin=Variable(torch.zeros(target.shape[0],11,target.shape[2],target.shape[3]).cuda().scatter_(1,target,1))
    smooth = 1.
    iflat = input.view(-1)
    tflat = target_bin.view(-1)
    intersection = (iflat * tflat).sum()
    return 1 - ((2. * intersection + smooth) /
            (iflat.sum() + tflat.sum() + smooth))

def weighted_loss(input,target,weight,size_average=True):
    n,c,h,w=input.size()
    target_bin=Variable(torch.zeros(n,c,h,w).cuda()).scatter_(1,target,1)
    target_bin=target_bin.transpose(1,2).transpose(2,3).contiguous().view(n*h*w,c).float()
    
    # NHWC
    input=F.softmax(input,dim=1).transpose(1,2).transpose(2,3).contiguous().view(n*h*w,c)
    input=input[target_bin>=0]
    input=input.view(-1,c)
    weight=weight.transpose(1,2).transpose(2,3).contiguous()
    weight=weight.view(n*h*w,1).repeat(1,c)
    '''
    mask=target>=0
    target=target[mask]
    target_bin=np.zeros((n*h*w,c),np.float)
    for i,term in enumerate(target):
        target_bin[i,int(term)]=1
    target_bin=torch.from_numpy(target_bin).float()
    target_bin=Variable(target_bin.cuda())
    '''
    loss=F.binary_cross_entropy(input,target_bin,weight=weight,size_average=False)
    if size_average:
        loss/=(target_bin>=0).data.sum().float()/c
    return loss

def bce2d_hed(input, target):
    """
    Binary Cross Entropy 2-Dimension loss
    """
    n, c, h, w = input.size()
    # assert(max(target) == 1)
    log_p = input.transpose(1, 2).transpose(2, 3).contiguous().view(1, -1)
    target_t = target.transpose(1, 2).transpose(2, 3).contiguous().view(1, -1).float().cuda()
    target_trans = target_t.clone()
    pos_index = (target_t >0)
    neg_index = (target_t ==0)
    target_trans[pos_index] = 1
    target_trans[neg_index] = 0
    pos_index = pos_index.data.cpu().numpy().astype(bool)
    neg_index = neg_index.data.cpu().numpy().astype(bool)
    weight = torch.Tensor(log_p.size()).fill_(0)
    weight = weight.numpy()
    pos_num = pos_index.sum()
    neg_num = neg_index.sum()
    sum_num = pos_num + neg_num
    weight[pos_index] = neg_num*1.0 / sum_num
    weight[neg_index] = pos_num*1.0 / sum_num

    weight = torch.from_numpy(weight)
    weight = weight.cuda()
    loss = F.binary_cross_entropy(log_p, target_t, weight, size_average=True)
    return loss

def bootstrapped_cross_entropy2d(input, target, K, weight=None, size_average=True):

    batch_size = input.size()[0]

    def _bootstrap_xentropy_single(input, target, K, weight=None, size_average=True):
        n, c, h, w = input.size()
        log_p = F.log_softmax(input, dim=1)
        log_p = log_p.transpose(1, 2).transpose(2, 3).contiguous().view(-1, c)
        log_p = log_p[target.view(n * h * w, 1).repeat(1, c) >= 0]
        log_p = log_p.view(-1, c)

        mask = target >= 0
        target = target[mask]
        loss = F.nll_loss(log_p, target, weight=weight, ignore_index=250,
                          reduce=False, size_average=False)
        topk_loss, _ = loss.topk(K)
        reduced_topk_loss = topk_loss.sum() / K

        return reduced_topk_loss

    loss = 0.0
    # Bootstrap from each image not entire batch
    for i in range(batch_size):
        loss += _bootstrap_xentropy_single(input=torch.unsqueeze(input[i], 0),
                                           target=torch.unsqueeze(target[i], 0),
                                           K=K,
                                           weight=weight,
                                           size_average=size_average)
    return loss / float(batch_size)

# another implimentation for dice loss
import torch
from torch.autograd import Function, Variable
class DiceCoeff(Function):
    """Dice coeff for individual examples"""
    def forward(self, input, target):
        self.save_for_backward(input, target)
        self.inter = torch.dot(input.view(-1), target.view(-1)) + 0.0001
        self.union = torch.sum(input) + torch.sum(target) + 0.0001
        t = 2 * self.inter.float() / self.union.float()
        return t
    # This function has only a single output, so it gets only one gradient
    def backward(self, grad_output):
        input, target = self.saved_variables
        grad_input = grad_target = None
        if self.needs_input_grad[0]:
            grad_input = grad_output * 2 * (target * self.union + self.inter) \
                         / self.union * self.union
        if self.needs_input_grad[1]:
            grad_target = None
        return grad_input, grad_target
def dice_coeff(input, target):
    target_bin=Variable(torch.zeros(1,11,target.shape[2],target.shape[3]).cuda().scatter_(1,target,1).float())
    """Dice coeff for batches"""
    if input.is_cuda:
        s = torch.FloatTensor(1).cuda().zero_()
    else:
        s = torch.FloatTensor(1).zero_()
    for i, c in enumerate(zip(input, target_bin)):
        s = s + DiceCoeff().forward(c[0], c[1])
    return s / (i + 1)