losses.py

import torch
import torch.nn.functional as F
import torch.nn as nn
from typing import Tuple, List
import numpy as np

#print(torch.__version__)

# KL divergence between the dirichlet distribution parameterized by alpha and the uniform dirichlet
# computes KLD per sample, expects input size of BxC
def KL(alpha, device):
    beta = torch.ones(1,alpha.size()[1]).to(device)
    S_alpha = torch.sum(alpha, 1, keepdim=True)
    S_beta = torch.sum(beta, 1, keepdim=True)
    #print(S_alpha)
    #print(S_beta)
    
    log_gamma_alpha = torch.lgamma(S_alpha) - torch.sum(torch.lgamma(alpha), 1, keepdim=True)
    log_gamma_beta = torch.lgamma(S_beta) - torch.sum(torch.lgamma(beta), 1, keepdim=True)
    #print(log_gamma_alpha)
    #print(log_gamma_beta)
    
    dg_a = torch.digamma(alpha)
    dg_a0 = torch.digamma(S_alpha)
    #print(dg_a)
    #print(dg_a0)
    
    kl = torch.sum((alpha-beta)*(dg_a-dg_a0) , 1, keepdim=True) + log_gamma_alpha - log_gamma_beta
    return kl

# MSE loss for evidential classification DNN
# First finds the expected probabilities for each class given dirichlet param. by alpha
# Then computes MSE and KLD components of the loss, expects alpha=BxC, p=BxC
def evidential_mse(alpha, p, device, weights=None):
    S = torch.sum(alpha, 1, keepdim=True)
    evidence = alpha - 1
    predicted_mass = alpha/S
    sq_err = (p-predicted_mass)**2
    if weights is not None:
        weighted_sq_err = sq_err * weights
    else:
        weighted_sq_err = sq_err
    #print(weighted_sq_err)
    sq_loss = torch.sum(weighted_sq_err, 1, keepdim=True)
    #print(sq_loss)
    dirichlet_variance = alpha*(S-alpha)/(S*S*(S+1))
    var = torch.sum(dirichlet_variance, 1, keepdim=True)
    #print(var)
    alpha_tilde = evidence*(1-p) + 1
    #print(alpha_tilde)
    kl = KL(alpha_tilde, device)
    return (sq_loss + var), kl

def evidential_ce(alpha, p, device):
    S = torch.sum(alpha, 1, keepdim=True)
    evidence = alpha - 1
    predicted_mass = alpha/S
    ce_loss = F.binary_cross_entropy(predicted_mass, p)
    #print(sq_loss)
    dirichlet_variance = alpha*(S-alpha)/(S*S*(S+1))
    var = torch.sum(dirichlet_variance, 1, keepdim=True)
    #print(var)
    alpha_tilde = evidence*(1-p) + 1
    #print(alpha_tilde)
    kl = KL(alpha_tilde, device)
    return (ce_loss + var), kl

# wrapper function for evidential_mse
# expects pred=BxC
def evidential_loss(pred, target, nclasses):
    alpha = F.softplus(pred)+1
    target_oh = F.one_hot(target, nclasses).float()
    mse, kl = evidential_ce(alpha, target_oh, alpha.device)
    loss = torch.mean(mse + 0.1*kl)
    return loss

# get the probability and evidential vacuity based on the network's logits
def evidential_prob_vacuity(logits, nclasses):
    alpha = F.softplus(logits)+1
    S = torch.sum(alpha, dim=1, keepdim=True)
    prob = alpha / S
    vacuity = nclasses / S
    return prob, vacuity

# CDF based probability dist. approximation from Zhibin
# given parameters to the laplacian distribution mu/sigma, get the discrete probabilities
# for a distribution between 0 and 1, of n bins. Expects mu/sigma Bx2
def laplace_cdf(mu_sigma, num_classes, device):
    mu = mu_sigma[:,0:1]
    sigma = mu_sigma[:,1:2]
    uv = torch.linspace(0, 1, steps=num_classes+1).to(device)
    u = uv[1:]
    v = uv[:-1]
    # print(mu)
    # print(sigma)
    # print(u)
    # print(v)
    epsilon = torch.tensor(1e-8).to(device)
    p = 0.5 * ( torch.sign(u-mu)*(1 - torch.exp(-torch.abs(u-mu) / sigma)) - torch.sign(v-mu)*(1 - torch.exp(-torch.abs(v-mu) / sigma)) ) + epsilon
    p = p / torch.sum(p, 1, keepdim=True)
    return p

def laplace_cdf_loss(pred, target, nclasses):
    pred_categorical = laplace_cdf(F.sigmoid(pred), nclasses, pred.device)
    target_oh = 0.9*F.one_hot(target, nclasses) + 0.1/nclasses
    loss = F.binary_cross_entropy(pred_categorical, target_oh)
    return loss

class SupConLoss(nn.Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR"""
    def __init__(self, temperature=0.07, contrast_mode='all',
                 base_temperature=0.07):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf
        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """
        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(device)
        else:
            mask = mask.float().to(device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)
        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # tile mask
        mask = mask.repeat(anchor_count, contrast_count)
        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
            0
        )
        mask = mask * logits_mask

        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, batch_size).mean()

        return loss
    
class CLIPLoss(torch.nn.Module):
  """
  Loss function for multimodal contrastive learning based off of the CLIP paper.
  
  Embeddings are taken, L2 normalized and dot product between modalities is calculated to generate a cosine
  similarity between all combinations of subjects in a cross-modal fashion. Tempered by temperature.
  Loss is calculated attempting to match each subject's embeddings between the modalities i.e. the diagonal. 
  """
  def __init__(self, 
               temperature: float,
               lambda_0: float = 0.5) -> None:
    super(CLIPLoss, self).__init__()

    self.temperature = temperature
    self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')

    if lambda_0 > 1 or lambda_0 < 0:
      raise ValueError('lambda_0 must be a float between 0 and 1.')
    self.lambda_0 = lambda_0
    self.lambda_1 = 1-lambda_0

  def forward(self, out0: torch.Tensor, out1: torch.Tensor, indices: List[int] = None) -> Tuple:
    # normalize the embedding onto the unit hypersphere
    out0 = nn.functional.normalize(out0, dim=1)
    out1 = nn.functional.normalize(out1, dim=1)

    #logits = torch.matmul(out0, out1.T) * torch.exp(torch.tensor(self.temperature))
    logits = torch.matmul(out0, out1.T) / self.temperature
    labels = torch.arange(len(out0), device=out0.device)

    #labels = labels.unsqueeze(1).expand(labels.shape[0], logits.shape[1])

    loss_0 = self.lambda_0 * self.cross_entropy(logits, labels)
    loss_1 = self.lambda_1 * self.cross_entropy(logits.T, labels)
    loss = loss_0 + loss_1
  
    return loss, logits, labels

def loss_coteaching(y_1, y_2, t, forget_rate, abs_loss=None):
    loss_1 = F.cross_entropy(y_1, t, reduce = False)
    ind_1_sorted = np.argsort(loss_1.cpu().data).cuda()
    loss_1_sorted = loss_1[ind_1_sorted]

    loss_2 = F.cross_entropy(y_2, t, reduce = False)
    ind_2_sorted = np.argsort(loss_2.cpu().data).cuda()

    remember_rate = 1 - forget_rate
    num_remember = int(remember_rate * len(loss_1_sorted))

    ind_1_update=ind_1_sorted[:num_remember]
    ind_2_update=ind_2_sorted[:num_remember]
    # exchange
    if abs_loss is None:
        loss_1_update = F.cross_entropy(y_1[ind_2_update], t[ind_2_update])
        loss_2_update = F.cross_entropy(y_2[ind_1_update], t[ind_1_update])
    else:    
        loss_1_update = abs_loss.compute(y_1[ind_2_update], t[ind_2_update])
        loss_2_update = abs_loss.compute(y_2[ind_1_update], t[ind_1_update])

    return torch.sum(loss_1_update), torch.sum(loss_2_update)

class CeLossAbstain(object):
    """
    Cross-entropy-like loss. Introduces a K+1-th class, abstention, which is a
    learned estimate of aleatoric uncertainty. When the network abstains,
    the loss will be computed with the ground truth, but, the network incurs
    loss for using the abstension
    """

    def __init__(self, loss_weight=1, ab_weight=0.3, reduction="sum", ab_logitpath="joined"):
        self.loss_weight = loss_weight
        self.reduction = reduction
        self.ab_weight = ab_weight
        self.ab_logitpath = ab_logitpath
        self.criterion = nn.NLLLoss(reduction=reduction)
        assert self.ab_logitpath == "joined" or self.ab_logitpath == "separate"
        print(
            f"setup CE Abstain Loss with loss_weight:{loss_weight}, "
            + f"ab_penalty:{ab_weight}, ab_path:{ab_logitpath} and reduction:{reduction}"
        )

    def to(self, device):
        self.criterion = self.criterion.to(device)

    def compute(self, logits, target):
        if self.loss_weight == 0:
            return torch.tensor(0, device=target.device)

        B, K_including_abs = logits.shape  # (B, K+1)
        K = K_including_abs - 1
        assert K >= 2, "CeLossAbstain input must have >= 2 classes not including abstention"

        # virtual_pred = (1-alpha) * pred + alpha * target
        if self.ab_logitpath == "joined":
            abs_pred = logits.softmax(dim=1)[:, K : K + 1]  # (B, 1)
        elif self.ab_logitpath == "separate":
            abs_pred = logits.sigmoid()[:, K : K + 1]  # (B, 1)
        class_pred = logits[:, :K].softmax(dim=1)  # (B, K)
        target_oh = nn.functional.one_hot(target, num_classes=K)
        virtual_pred = (1 - abs_pred) * class_pred + abs_pred * target_oh  # (B, K)

        loss_pred = self.criterion(torch.log(virtual_pred), target)
        loss_abs = -torch.log(1 - abs_pred).squeeze()  # (B)

        if self.reduction == "mean":
            loss_abs = loss_abs.mean()
        elif self.reduction == "sum":
            loss_abs = loss_abs.sum()

        return self.loss_weight * (loss_pred + self.ab_weight * loss_abs)