123456789101112131415161718192021222324252627 |
- #! /usr/bin/python
- # -*- encoding: utf-8 -*-
- 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, nOut, nClasses, **kwargs):
- super(LossFunction, self).__init__()
- self.test_normalize = True
-
- self.criterion = torch.nn.CrossEntropyLoss()
- self.fc = nn.Linear(nOut,nClasses)
- print('Initialised Softmax Loss')
- def forward(self, x, label=None):
- x = self.fc(x)
- nloss = self.criterion(x, label)
- prec1 = accuracy(x.detach(), label.detach(), topk=(1,))[0]
- return nloss, prec1
|