angleproto.py 1.2 KB

123456789101112131415161718192021222324252627282930313233343536373839
  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 accuracy
  8. class LossFunction(nn.Module):
  9. def __init__(self, init_w=10.0, init_b=-5.0, **kwargs):
  10. super(LossFunction, self).__init__()
  11. self.test_normalize = True
  12. self.w = nn.Parameter(torch.tensor(init_w))
  13. self.b = nn.Parameter(torch.tensor(init_b))
  14. self.criterion = torch.nn.CrossEntropyLoss()
  15. print('Initialised AngleProto')
  16. def forward(self, x, label=None):
  17. assert x.size()[1] >= 2
  18. out_anchor = torch.mean(x[:,1:,:],1)
  19. out_positive = x[:,0,:]
  20. stepsize = out_anchor.size()[0]
  21. cos_sim_matrix = F.cosine_similarity(out_positive.unsqueeze(-1),out_anchor.unsqueeze(-1).transpose(0,2))
  22. torch.clamp(self.w, 1e-6)
  23. cos_sim_matrix = cos_sim_matrix * self.w + self.b
  24. label = torch.from_numpy(numpy.asarray(range(0,stepsize))).cuda()
  25. nloss = self.criterion(cos_sim_matrix, label)
  26. prec1 = accuracy(cos_sim_matrix.detach(), label.detach(), topk=(1,))[0]
  27. return nloss, prec1