"""
Loss function classes for safe saving
"""
from typing import Callable, Any
import torch
from torch import nn, Tensor
[docs]
class BaseLoss(nn.Module):
"""
A base class for loss functions.
This class is used to add all loss function classes to the list of safe PyTorch classes when
loading saved architectures.
"""
def __init__(self, loss_func: type, *args: Any, **kwargs: Any) -> None:
"""
Parameters
----------
loss_func : type
Loss function class to be used
*args
Optional arguments to be passed to loss_func
**kwargs
Optional keyword arguments to be passed to loss_func
"""
super().__init__()
self._args: tuple[Any, ...] = args
self._kwargs: dict[str, Any] = kwargs
self._loss_func: Callable = loss_func(*self._args, **self._kwargs)
# Adds all loss classes to list of safe PyTorch classes when loading saved architectures
torch.serialization.add_safe_globals([self.__class__])
def __repr__(self) -> str:
"""
Returns a string representation of the loss function.
Returns
-------
str
String representation of the loss function
"""
key: str
args: str
extra_repr: str
val: Any
args = ', '.join(repr(val) for val in self._args)
args += ', '.join(f'{key}={repr(val)}' for key, val in self._kwargs.items())
args += (', ' if args else '') + extra_repr if (extra_repr := self.extra_repr()) else ''
return f'{self.__class__.__name__}({args})'
def __getstate__(self) -> dict[str, Any]:
"""
Returns a dictionary containing the state of the loss function for pickling.
Returns
-------
dict[str, Any]
Dictionary containing the state of the loss function
"""
return {'args': self._args, 'kwargs': self._kwargs}
def __setstate__(self, state: dict[str, Any]) -> None:
"""
Sets the state of the loss function for pickling.
Parameters
----------
state : dict[str, Any]
Dictionary containing the state of the loss function
"""
super().__init__()
self._args = state['args']
self._kwargs = state['kwargs']
[docs]
def forward(self, output: Tensor, target: Tensor) -> Tensor:
"""
Forward pass of the loss function.
Parameters
----------
output : Tensor
Output from the network with shape (N,...), where N is the number of elements
N predictions from the network
target : Tensor
Target values with shape (N,...)
Returns
-------
Tensor
Loss value with shape (1)
"""
return self._loss_func(output, target)
[docs]
class GeneralLoss(BaseLoss):
"""
A general loss function that can be used with any torch.nn loss function class.
"""
def __init__(self, loss_func: str, *args: Any, **kwargs: Any) -> None:
"""
Parameters
----------
loss_func : str
Loss function name from torch.nn
*args
Optional arguments to be passed to loss_func
**kwargs
Optional keyword arguments to be passed to loss_func
"""
self._name: str = loss_func
super().__init__(getattr(nn, self._name), *args, **kwargs)
def __getstate__(self) -> dict[str, Any]:
return super().__getstate__() | {'name': self._name}
def __setstate__(self, state: dict[str, Any]) -> None:
super().__setstate__(state)
self._name = state.get('name', type(self).__name__)
self._loss_func = getattr(nn, self._name)(*self._args, **self._kwargs)
[docs]
class CrossEntropyLoss(GeneralLoss):
"""
Cross entropy loss function.
"""
def __init__(self, *args: Any, **kwargs: Any) -> None:
"""
Parameters
----------
*args
Optional arguments to be passed to CrossEntropyLoss
**kwargs
Optional keyword arguments to be passed to CrossEntropyLoss
"""
super().__init__('CrossEntropyLoss', *args, **kwargs)
[docs]
class GaussianNLLLoss(GeneralLoss):
"""
Gaussian negative log likelihood loss function.
"""
def __init__(self, *args: Any, dim: int = 1, **kwargs: Any) -> None:
"""
Parameters
----------
dim : int, Optional
Dimension of the target that contains the mean and uncertainty, default = 1
*args
Optional arguments to be passed to GaussianNLLLoss
**kwargs
Optional keyword arguments to be passed to GaussianNLLLoss
"""
super().__init__('GaussianNLLLoss', *args, **kwargs)
self._dim: int = dim
def __getstate__(self) -> dict[str, Any]:
return super().__getstate__() | {'dim': self._dim}
def __setstate__(self, state: dict[str, Any]) -> None:
super().__setstate__(state)
self._dim = state['dim']
[docs]
def forward(self, output: Tensor, target: Tensor) -> Tensor:
if target.shape[self._dim] != 2 or output.shape[self._dim] != 1:
raise ValueError(f'Target shape must be 2 and output shape must be 1 at dim '
f'{self._dim}, got target shape {target.shape} and output shape '
f'{output.shape}')
idx: list[list[int] | slice] = [slice(None)] * target.dim()
idx[self._dim] = [0, 1]
return self._loss_func(output[*idx][0], target[*idx][0], target[*idx][1] ** 2)
[docs]
class MSELoss(GeneralLoss):
"""
Mean Squared Error (MSE) loss function.
"""
def __init__(self, *args: Any, **kwargs: Any) -> None:
"""
Parameters
----------
*args
Optional arguments to be passed to MSELoss
**kwargs
Optional keyword arguments to be passed to MSELoss
"""
super().__init__('MSELoss', *args, **kwargs)
__all__ = ['BaseLoss', 'CrossEntropyLoss', 'GaussianNLLLoss', 'GeneralLoss', 'MSELoss']