softmaxproto.py 722 B

1234567891011121314151617181920212223242526272829
  1. #! /usr/bin/python
  2. # -*- encoding: utf-8 -*-
  3. import torch
  4. import torch.nn as nn
  5. import loss.softmax as softmax
  6. import loss.angleproto as angleproto
  7. class LossFunction(nn.Module):
  8. def __init__(self, **kwargs):
  9. super(LossFunction, self).__init__()
  10. self.test_normalize = True
  11. self.softmax = softmax.LossFunction(**kwargs)
  12. self.angleproto = angleproto.LossFunction(**kwargs)
  13. print('Initialised SoftmaxPrototypical Loss')
  14. def forward(self, x, label=None):
  15. assert x.size()[1] == 2
  16. nlossS, prec1 = self.softmax(x.reshape(-1,x.size()[-1]), label.repeat_interleave(2))
  17. nlossP, _ = self.angleproto(x,None)
  18. return nlossS+nlossP, prec1