loss.py

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



def cosine_sim(im, s):
    """Cosine similarity between all the image and sentence pairs
    """
    return im.mm(s.t())


def order_sim(im, s):
    """Order embeddings similarity measure $max(0, s-im)$
    """
    YmX = (s.unsqueeze(1).expand(s.size(0), im.size(0), s.size(1))
           - im.unsqueeze(0).expand(s.size(0), im.size(0), s.size(1)))
    score = -YmX.clamp(min=0).pow(2).sum(2).sqrt().t()
    return score


def euclidean_sim(im, s):
    """L2 distance
    """
    YmX = (s.unsqueeze(1).expand(s.size(0), im.size(0), s.size(1))
           - im.unsqueeze(0).expand(s.size(0), im.size(0), s.size(1)))
    score = -YmX.pow(2).sum(2).t()
    return score


def L1_sim(im, s):
    """L1 distance
    """
    YmX = (s.unsqueeze(1).expand(s.size(0), im.size(0), s.size(1))
           - im.unsqueeze(0).expand(s.size(0), im.size(0), s.size(1)))
    score = -YmX.abs().sum(2).t()
    return score

def L1_sim_norm(im, s):
    """L1 normalization distance  1 - L_1/K
    """
    YmX = (s.unsqueeze(1).expand(s.size(0), im.size(0), s.size(1))
           - im.unsqueeze(0).expand(s.size(0), im.size(0), s.size(1)))
    score = YmX.abs().sum(2).t()/im.size(1) -1
    return score


def L2_sim(im, s):
    """L2 distance
    """
    YmX = (s.unsqueeze(1).expand(s.size(0), im.size(0), s.size(1))
           - im.unsqueeze(0).expand(s.size(0), im.size(0), s.size(1)))
    score = -YmX.pow(2).sum(2).t()
    return score

def L2_sim_norm(im, s):
    """L2 normalization distance  1 - L_2/K
    """
    YmX = (s.unsqueeze(1).expand(s.size(0), im.size(0), s.size(1))
           - im.unsqueeze(0).expand(s.size(0), im.size(0), s.size(1)))
    score = YmX.pow(2).sum(2).t()/im.size(1) - 1
    return score


def jaccard_sim(im, s):
    im_bs = im.size(0)
    s_bs = s.size(0)
    im = im.unsqueeze(1).expand(-1,s_bs,-1)
    s = s.unsqueeze(0).expand(im_bs,-1,-1)
    intersection = torch.min(im,s).sum(-1)
    union = torch.max(im,s).sum(-1)
    score = intersection / union
    return score


NAME_TO_SIM = {'cosine': cosine_sim, 'order': order_sim, 'euclidean': euclidean_sim, 'jaccard': jaccard_sim}

def get_sim(name):
    assert name in NAME_TO_SIM, '%s not supported.'%name
    return NAME_TO_SIM[name]


class TripletLoss(nn.Module):
    """
    triplet ranking loss
    """

    def __init__(self, margin=0, measure=False, max_violation=False, cost_style='sum', direction='all'):
        super(TripletLoss, self).__init__()
        self.margin = margin
        self.cost_style = cost_style
        self.direction = direction
        if measure == 'order':
            self.sim = order_sim
        elif measure == 'euclidean':
            self.sim = euclidean_sim
        elif measure == 'jaccard':
            self.sim = jaccard_sim
        elif measure == 'l1':
            self.sim = L1_sim
        elif measure == 'l2':
            self.sim = L2_sim
        elif measure == 'l1_norm':
            self.sim = L1_sim_norm
        elif measure == 'l2_norm':
            self.sim = L2_sim_norm
        else:
            self.sim = cosine_sim

        self.max_violation = max_violation

    def forward(self, s, im):
        # compute video-sentence score matrix
        scores = self.sim(im, s)
        diagonal = scores.diag().view(im.size(0), 1)
        d1 = diagonal.expand_as(scores)
        d2 = diagonal.t().expand_as(scores)

        # clear diagonals
        mask = torch.eye(scores.size(0)) > .5
        I = Variable(mask)
        if torch.cuda.is_available():
            I = I.cuda()

        cost_s = None
        cost_im = None
        # compare every diagonal score to scores in its column
        if self.direction in ['v2t', 'all']:
            # caption retrieval
            cost_s = (self.margin + scores - d1).clamp(min=0)
            cost_s = cost_s.masked_fill_(I, 0)
        # compare every diagonal score to scores in its row
        if self.direction in ['t2v', 'all']:
            # video retrieval
            cost_im = (self.margin + scores - d2).clamp(min=0)
            cost_im = cost_im.masked_fill_(I, 0)

        # keep the maximum violating negative for each query
        if self.max_violation:
            if cost_s is not None:
                cost_s = cost_s.max(1)[0]
            if cost_im is not None:
                cost_im = cost_im.max(0)[0]

        if cost_s is None:
            cost_s = Variable(torch.zeros(1)).cuda()
        if cost_im is None:
            cost_im = Variable(torch.zeros(1)).cuda()

        if self.cost_style == 'sum':
            return cost_s.sum() + cost_im.sum()
        else:
            return cost_s.mean() + cost_im.mean()




class DtlLoss(nn.Module):
    def __init__(self):
        super(DtlLoss, self).__init__()
        self.temp = 0.07

    def forward(self, feat_targets, feat, feat_v):

        sim_v2t_m = cosine_sim(feat_v, feat_targets) / self.temp
        sim_t2v_m = cosine_sim(feat_targets, feat_v) / self.temp

        sim_v2t_targets = F.softmax(sim_v2t_m, dim=1)
        sim_t2v_targets = F.softmax(sim_t2v_m, dim=1)

        sim_v2t = cosine_sim(feat_v, feat) / self.temp
        sim_t2v = cosine_sim(feat, feat_v) / self.temp

        loss_v2t = -torch.sum(F.log_softmax(sim_v2t, dim=1) * sim_v2t_targets, dim=1).mean()
        loss_t2v = -torch.sum(F.log_softmax(sim_t2v, dim=1) * sim_t2v_targets, dim=1).mean()

        loss_ita = (loss_v2t + loss_t2v) / 2

        return loss_ita




class dtl_feat(nn.Module):

    def __init__(self):
        super(dtl_feat, self).__init__()
        self.beta = 4
        self.smooth_l1 = nn.SmoothL1Loss(size_average=None, reduce=None, reduction='mean', beta=self.beta)
    def forward(self, feat_target, feat):

        loss = self.smooth_l1(feat.float(), feat_target.float())
        return loss