| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485 | #! /usr/bin/python# -*- encoding: utf-8 -*-import torchimport torch.nn as nnimport torch.nn.functional as Fimport time, pdb, numpyfrom utils import tuneThresholdfromScoreimport randomclass LossFunction(nn.Module):    def __init__(self, hard_rank=0, hard_prob=0, margin=0, **kwargs):        super(LossFunction, self).__init__()        self.test_normalize = True                self.hard_rank  = hard_rank        self.hard_prob  = hard_prob        self.margin     = margin        print('Initialised Triplet Loss')    def forward(self, x, label=None):        assert x.size()[1] == 2                out_anchor      = F.normalize(x[:,0,:], p=2, dim=1)        out_positive    = F.normalize(x[:,1,:], p=2, dim=1)        stepsize        = out_anchor.size()[0]        output      = -1 * (F.pairwise_distance(out_anchor.unsqueeze(-1),out_positive.unsqueeze(-1).transpose(0,2))**2)        negidx      = self.mineHardNegative(output.detach())        out_negative = out_positive[negidx,:]        labelnp     = numpy.array([1]*len(out_positive)+[0]*len(out_negative))        ## calculate distances        pos_dist    = F.pairwise_distance(out_anchor,out_positive)        neg_dist    = F.pairwise_distance(out_anchor,out_negative)        ## loss function        nloss   = torch.mean(F.relu(torch.pow(pos_dist, 2) - torch.pow(neg_dist, 2) + self.margin))        scores = -1 * torch.cat([pos_dist,neg_dist],dim=0).detach().cpu().numpy()        return nloss, nloss    ## ===== ===== ===== ===== ===== ===== ===== =====    ## Hard negative mining    ## ===== ===== ===== ===== ===== ===== ===== =====    def mineHardNegative(self, output):        negidx = []        for idx, similarity in enumerate(output):            simval, simidx = torch.sort(similarity,descending=True)            if self.hard_rank < 0:                ## Semi hard negative mining                semihardidx = simidx[(similarity[idx] - self.margin < simval) &  (simval < similarity[idx])]                if len(semihardidx) == 0:                    negidx.append(random.choice(simidx))                else:                    negidx.append(random.choice(semihardidx))            else:                ## Rank based negative mining                                simidx = simidx[simidx!=idx]                if random.random() < self.hard_prob:                    negidx.append(simidx[random.randint(0, self.hard_rank)])                else:                    negidx.append(random.choice(simidx))        return negidx
 |