softmax.py 633 B

123456789101112131415161718192021222324252627
  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, nOut, nClasses, **kwargs):
  10. super(LossFunction, self).__init__()
  11. self.test_normalize = True
  12. self.criterion = torch.nn.CrossEntropyLoss()
  13. self.fc = nn.Linear(nOut,nClasses)
  14. print('Initialised Softmax Loss')
  15. def forward(self, x, label=None):
  16. x = self.fc(x)
  17. nloss = self.criterion(x, label)
  18. prec1 = accuracy(x.detach(), label.detach(), topk=(1,))[0]
  19. return nloss, prec1