93 lines
3.2 KiB
Python
93 lines
3.2 KiB
Python
from .module import Module
|
|
from norch.autograd.functions import *
|
|
import norch
|
|
from abc import ABC
|
|
|
|
class Loss(Module, ABC):
|
|
"Abstract class for loss functions"
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, predictions, labels):
|
|
raise NotImplementedError
|
|
|
|
def __call__(self, *inputs):
|
|
return self.forward(*inputs)
|
|
|
|
|
|
class MSELoss(Loss):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, predictions, target):
|
|
assert target.shape == predictions.shape, \
|
|
"Labels and predictions shape does not match: {} and {}".format(target.shape, predictions.shape)
|
|
|
|
cost = ((predictions - target) ** 2).sum() / predictions.numel
|
|
return cost
|
|
|
|
|
|
class CrossEntropyLoss(Loss):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, input, target):
|
|
assert isinstance(input, norch.Tensor), \
|
|
"Cross entropy argument 'input' must be Tensor, not {}".format(type(input))
|
|
|
|
assert isinstance(target, norch.Tensor), \
|
|
"Cross entropy argument 'target' must be Tensor, not {}".format(type(target))
|
|
if input.ndim > 2:
|
|
input = input.squeeze(-1)
|
|
|
|
if input.ndim == 1:
|
|
if target.numel == 1:
|
|
num_classes = input.shape[0]
|
|
target = norch.one_hot_encode(target, num_classes).to(target.device)
|
|
|
|
logits = norch.softmax(input, dim=0)
|
|
target = target.reshape(logits.shape)
|
|
cost = -(logits.log() * target).sum()
|
|
|
|
else:
|
|
# target -> class probabilities (one-hot encoded)
|
|
assert target.shape == input.shape, \
|
|
"Input and target shape does not match: {} and {}".format(input.shape, target.shape)
|
|
logits = norch.softmax(input, dim=0)
|
|
target = target.reshape(logits.shape)
|
|
cost = -(logits.log() * target).sum()
|
|
|
|
|
|
elif input.ndim == 2:
|
|
if target.ndim > 1:
|
|
target = target.squeeze(-1)
|
|
# batched
|
|
if target.ndim == 1:
|
|
# target -> Ground truth class indices:
|
|
num_classes = input.shape[1]
|
|
|
|
target = norch.one_hot_encode(target, num_classes).to(target.device)
|
|
|
|
batch_size = input.shape[0]
|
|
logits = norch.softmax(input, dim=1)
|
|
target = target.reshape(logits.shape)
|
|
cost = -(logits.log() * target).sum() / batch_size
|
|
|
|
else:
|
|
# target -> class probabilities (one-hot encoded)
|
|
assert target.shape == input.shape, \
|
|
"Input and target shape does not match: {} and {}".format(input.shape, target.shape)
|
|
|
|
batch_size = input.shape[0]
|
|
logits = norch.softmax(input, dim=1)
|
|
target = target.reshape(logits.shape)
|
|
cost = -(logits.log() * target).sum() / batch_size
|
|
|
|
if input.requires_grad:
|
|
cost.grad_fn = CrossEntropyLossBackward(input, target)
|
|
|
|
return cost
|
|
|
|
|
|
|