aamsoftmax.py 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. #! /usr/bin/python
  2. # -*- encoding: utf-8 -*-
  3. # Adapted from https://github.com/wujiyang/Face_Pytorch (Apache License)
  4. import torch
  5. import torch.nn as nn
  6. import torch.nn.functional as F
  7. import time, pdb, numpy, math
  8. from utils import accuracy
  9. class LossFunction(nn.Module):
  10. def __init__(self, nOut, nClasses, margin=0.3, scale=15, easy_margin=False, **kwargs):
  11. super(LossFunction, self).__init__()
  12. self.test_normalize = True
  13. self.m = margin
  14. self.s = scale
  15. self.in_feats = nOut
  16. self.weight = torch.nn.Parameter(torch.FloatTensor(nClasses, nOut), requires_grad=True)
  17. self.ce = nn.CrossEntropyLoss()
  18. nn.init.xavier_normal_(self.weight, gain=1)
  19. self.easy_margin = easy_margin
  20. self.cos_m = math.cos(self.m)
  21. self.sin_m = math.sin(self.m)
  22. # make the function cos(theta+m) monotonic decreasing while theta in [0°,180°]
  23. self.th = math.cos(math.pi - self.m)
  24. self.mm = math.sin(math.pi - self.m) * self.m
  25. print('Initialised AAMSoftmax margin %.3f scale %.3f'%(self.m,self.s))
  26. def forward(self, x, label=None):
  27. assert x.size()[0] == label.size()[0]
  28. assert x.size()[1] == self.in_feats
  29. # cos(theta)
  30. cosine = F.linear(F.normalize(x), F.normalize(self.weight))
  31. # cos(theta + m)
  32. sine = torch.sqrt((1.0 - torch.mul(cosine, cosine)).clamp(0, 1))
  33. phi = cosine * self.cos_m - sine * self.sin_m
  34. if self.easy_margin:
  35. phi = torch.where(cosine > 0, phi, cosine)
  36. else:
  37. phi = torch.where((cosine - self.th) > 0, phi, cosine - self.mm)
  38. #one_hot = torch.zeros(cosine.size(), device='cuda' if torch.cuda.is_available() else 'cpu')
  39. one_hot = torch.zeros_like(cosine)
  40. one_hot.scatter_(1, label.view(-1, 1), 1)
  41. output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
  42. output = output * self.s
  43. loss = self.ce(output, label)
  44. prec1 = accuracy(output.detach(), label.detach(), topk=(1,))[0]
  45. return loss, prec1