ge2e.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. #! /usr/bin/python
  2. # -*- encoding: utf-8 -*-
  3. ## Fast re-implementation of the GE2E loss (https://arxiv.org/abs/1710.10467)
  4. ## Numerically checked against https://github.com/cvqluu/GE2E-Loss
  5. import torch
  6. import torch.nn as nn
  7. import torch.nn.functional as F
  8. import time, pdb, numpy
  9. from utils import accuracy
  10. class LossFunction(nn.Module):
  11. def __init__(self, init_w=10.0, init_b=-5.0, **kwargs):
  12. super(LossFunction, self).__init__()
  13. self.test_normalize = True
  14. self.w = nn.Parameter(torch.tensor(init_w))
  15. self.b = nn.Parameter(torch.tensor(init_b))
  16. self.criterion = torch.nn.CrossEntropyLoss()
  17. print('Initialised GE2E')
  18. def forward(self, x, label=None):
  19. assert x.size()[1] >= 2
  20. gsize = x.size()[1]
  21. centroids = torch.mean(x, 1)
  22. stepsize = x.size()[0]
  23. cos_sim_matrix = []
  24. for ii in range(0,gsize):
  25. idx = [*range(0,gsize)]
  26. idx.remove(ii)
  27. exc_centroids = torch.mean(x[:,idx,:], 1)
  28. cos_sim_diag = F.cosine_similarity(x[:,ii,:],exc_centroids)
  29. cos_sim = F.cosine_similarity(x[:,ii,:].unsqueeze(-1),centroids.unsqueeze(-1).transpose(0,2))
  30. cos_sim[range(0,stepsize),range(0,stepsize)] = cos_sim_diag
  31. cos_sim_matrix.append(torch.clamp(cos_sim,1e-6))
  32. cos_sim_matrix = torch.stack(cos_sim_matrix,dim=1)
  33. torch.clamp(self.w, 1e-6)
  34. cos_sim_matrix = cos_sim_matrix * self.w + self.b
  35. label = torch.from_numpy(numpy.asarray(range(0,stepsize))).cuda()
  36. nloss = self.criterion(cos_sim_matrix.view(-1,stepsize), torch.repeat_interleave(label,repeats=gsize,dim=0).cuda())
  37. prec1 = accuracy(cos_sim_matrix.view(-1,stepsize).detach(), torch.repeat_interleave(label,repeats=gsize,dim=0).detach(), topk=(1,))[0]
  38. return nloss, prec1