triplet.py 2.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. #! /usr/bin/python
  2. # -*- encoding: utf-8 -*-
  3. import torch
  4. import torch.nn as nn
  5. import torch.nn.functional as F
  6. import time, pdb, numpy
  7. from utils import tuneThresholdfromScore
  8. import random
  9. class LossFunction(nn.Module):
  10. def __init__(self, hard_rank=0, hard_prob=0, margin=0, **kwargs):
  11. super(LossFunction, self).__init__()
  12. self.test_normalize = True
  13. self.hard_rank = hard_rank
  14. self.hard_prob = hard_prob
  15. self.margin = margin
  16. print('Initialised Triplet Loss')
  17. def forward(self, x, label=None):
  18. assert x.size()[1] == 2
  19. out_anchor = F.normalize(x[:,0,:], p=2, dim=1)
  20. out_positive = F.normalize(x[:,1,:], p=2, dim=1)
  21. stepsize = out_anchor.size()[0]
  22. output = -1 * (F.pairwise_distance(out_anchor.unsqueeze(-1),out_positive.unsqueeze(-1).transpose(0,2))**2)
  23. negidx = self.mineHardNegative(output.detach())
  24. out_negative = out_positive[negidx,:]
  25. labelnp = numpy.array([1]*len(out_positive)+[0]*len(out_negative))
  26. ## calculate distances
  27. pos_dist = F.pairwise_distance(out_anchor,out_positive)
  28. neg_dist = F.pairwise_distance(out_anchor,out_negative)
  29. ## loss function
  30. nloss = torch.mean(F.relu(torch.pow(pos_dist, 2) - torch.pow(neg_dist, 2) + self.margin))
  31. scores = -1 * torch.cat([pos_dist,neg_dist],dim=0).detach().cpu().numpy()
  32. return nloss, nloss
  33. ## ===== ===== ===== ===== ===== ===== ===== =====
  34. ## Hard negative mining
  35. ## ===== ===== ===== ===== ===== ===== ===== =====
  36. def mineHardNegative(self, output):
  37. negidx = []
  38. for idx, similarity in enumerate(output):
  39. simval, simidx = torch.sort(similarity,descending=True)
  40. if self.hard_rank < 0:
  41. ## Semi hard negative mining
  42. semihardidx = simidx[(similarity[idx] - self.margin < simval) & (simval < similarity[idx])]
  43. if len(semihardidx) == 0:
  44. negidx.append(random.choice(simidx))
  45. else:
  46. negidx.append(random.choice(semihardidx))
  47. else:
  48. ## Rank based negative mining
  49. simidx = simidx[simidx!=idx]
  50. if random.random() < self.hard_prob:
  51. negidx.append(simidx[random.randint(0, self.hard_rank)])
  52. else:
  53. negidx.append(random.choice(simidx))
  54. return negidx