123456789101112131415161718192021222324252627282930313233343536 |
- #! /usr/bin/python
- # -*- encoding: utf-8 -*-
- ## Re-implementation of prototypical networks (https://arxiv.org/abs/1703.05175).
- ## Numerically checked against https://github.com/cyvius96/prototypical-network-pytorch
- import torch
- import torch.nn as nn
- import torch.nn.functional as F
- import time, pdb, numpy
- from utils import accuracy
- class LossFunction(nn.Module):
- def __init__(self, **kwargs):
- super(LossFunction, self).__init__()
- self.test_normalize = False
- self.criterion = torch.nn.CrossEntropyLoss()
- print('Initialised Prototypical Loss')
- def forward(self, x, label=None):
- assert x.size()[1] >= 2
-
- out_anchor = torch.mean(x[:,1:,:],1)
- out_positive = x[:,0,:]
- stepsize = out_anchor.size()[0]
- output = -1 * (F.pairwise_distance(out_positive.unsqueeze(-1),out_anchor.unsqueeze(-1).transpose(0,2))**2)
- label = torch.from_numpy(numpy.asarray(range(0,stepsize))).cuda()
- nloss = self.criterion(output, label)
- prec1 = accuracy(output.detach(), label.detach(), topk=(1,))[0]
- return nloss, prec1
|