proto.py 1.1 KB

123456789101112131415161718192021222324252627282930313233343536
  1. #! /usr/bin/python
  2. # -*- encoding: utf-8 -*-
  3. ## Re-implementation of prototypical networks (https://arxiv.org/abs/1703.05175).
  4. ## Numerically checked against https://github.com/cyvius96/prototypical-network-pytorch
  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, **kwargs):
  12. super(LossFunction, self).__init__()
  13. self.test_normalize = False
  14. self.criterion = torch.nn.CrossEntropyLoss()
  15. print('Initialised Prototypical Loss')
  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. output = -1 * (F.pairwise_distance(out_positive.unsqueeze(-1),out_anchor.unsqueeze(-1).transpose(0,2))**2)
  22. label = torch.from_numpy(numpy.asarray(range(0,stepsize))).cuda()
  23. nloss = self.criterion(output, label)
  24. prec1 = accuracy(output.detach(), label.detach(), topk=(1,))[0]
  25. return nloss, prec1