amsoftmax.py 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. #! /usr/bin/python
  2. # -*- encoding: utf-8 -*-
  3. # Adapted from https://github.com/CoinCheung/pytorch-loss (MIT License)
  4. import torch
  5. import torch.nn as nn
  6. import torch.nn.functional as F
  7. import time, pdb, numpy
  8. from utils import accuracy
  9. class LossFunction(nn.Module):
  10. def __init__(self, nOut, nClasses, margin=0.3, scale=15, **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.W = torch.nn.Parameter(torch.randn(nOut, nClasses), requires_grad=True)
  17. self.ce = nn.CrossEntropyLoss()
  18. nn.init.xavier_normal_(self.W, gain=1)
  19. print('Initialised AMSoftmax m=%.3f s=%.3f'%(self.m,self.s))
  20. def forward(self, x, label=None):
  21. assert x.size()[0] == label.size()[0]
  22. assert x.size()[1] == self.in_feats
  23. x_norm = torch.norm(x, p=2, dim=1, keepdim=True).clamp(min=1e-12)
  24. x_norm = torch.div(x, x_norm)
  25. w_norm = torch.norm(self.W, p=2, dim=0, keepdim=True).clamp(min=1e-12)
  26. w_norm = torch.div(self.W, w_norm)
  27. costh = torch.mm(x_norm, w_norm)
  28. label_view = label.view(-1, 1)
  29. if label_view.is_cuda: label_view = label_view.cpu()
  30. delt_costh = torch.zeros(costh.size()).scatter_(1, label_view, self.m)
  31. if x.is_cuda: delt_costh = delt_costh.cuda()
  32. costh_m = costh - delt_costh
  33. costh_m_s = self.s * costh_m
  34. loss = self.ce(costh_m_s, label)
  35. prec1 = accuracy(costh_m_s.detach(), label.detach(), topk=(1,))[0]
  36. return loss, prec1