netloader 3.11.0__py3-none-any.whl

netloader/loss_funcs.py

@@ -0,0 +1,201 @@
+"""
+Loss function classes for safe saving
+"""
+from typing import Callable, Any
+
+import torch
+from torch import nn, Tensor
+
+
+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']
+
+    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)
+
+
+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['name']
+        self._loss_func = getattr(nn, self._name)(*self._args, **self._kwargs)
+
+    def extra_repr(self) -> str:
+        return f'loss_func={repr(self._name)}'
+
+
+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)
+
+
+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']
+
+    def extra_repr(self) -> str:
+        return f'dim={self._dim}'
+
+    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)
+
+
+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']

netloader/models/__init__.py

@@ -0,0 +1,22 @@
+"""
+Collects all classes in the models module
+"""
+from netloader.models.misc import MultiNetwork
+from netloader.models.convnext import (
+    ConvNeXt,
+    ConvNeXtTiny,
+    ConvNeXtBase,
+    ConvNeXtSmall,
+    ConvNeXtLarge,
+    ConvNeXtXLarge,
+)
+
+__all__ = [
+    'ConvNeXt',
+    'ConvNeXtTiny',
+    'ConvNeXtBase',
+    'MultiNetwork',
+    'ConvNeXtSmall',
+    'ConvNeXtLarge',
+    'ConvNeXtXLarge',
+]