netloader 3.11.0__py3-none-any.whl

netloader/networks/__init__.py

@@ -0,0 +1,48 @@
+"""
+Collects all architectures
+"""
+from warnings import warn
+
+from netloader.architectures import (
+    BaseArchitecture,
+    UtilityMixin,
+    Autoencoder,
+    Decoder,
+    Encoder,
+    load_arch,
+)
+
+BaseNetwork = BaseArchitecture
+load_net = load_arch
+
+
+warn(
+    'Importing architectures from netloader.networks is deprecated, please use '
+    'netloader.architectures',
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+__all__ = [
+    'UtilityMixin',
+    'BaseNetwork',
+    'Autoencoder',
+    'Decoder',
+    'Encoder',
+    'load_net',
+]
+
+_optional_imports: dict[str, set[str]] = {'flow': {'NormFlow', 'NormFlowEncoder'}}
+
+try:
+    from netloader.architectures.flows import NormFlow, NormFlowEncoder
+    __all__.extend(['NormFlow', 'NormFlowEncoder'])
+except (ModuleNotFoundError, ImportError):
+    pass
+
+
+def __getattr__(name):
+    if name not in __all__ and name in _optional_imports['flow']:
+        raise ImportError(f'Cannot import {name}, normalising flow architectures require the '
+                          f'package Zuko')
+    raise AttributeError(f"module {__name__} has no attribute {name}")

netloader/schedulers.py

@@ -0,0 +1,409 @@
+"""
+Scheduler classes.
+"""
+from typing import Any
+from abc import ABC, abstractmethod
+
+import torch
+import numpy as np
+
+
+class BaseScheduler(ABC):
+    """
+    Abstract base class for schedulers.
+    """
+    def __init__(self) -> None:
+        # Adds all scheduler classes to list of safe PyTorch classes when loading saved architectures
+        torch.serialization.add_safe_globals([self.__class__])
+
+    @abstractmethod
+    def __call__(self) -> float:
+        """
+        Returns the current value of the scheduler.
+        """
+
+    @abstractmethod
+    def __getstate__(self) -> dict[str, Any]:
+        """
+        Returns the state of the scheduler.
+
+        Returns
+        -------
+        dict[str, Any]
+            State of the scheduler
+        """
+
+    @abstractmethod
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """
+        Sets the state of the scheduler.
+
+        Parameters
+        ----------
+        state : dict[str, Any]
+            State of the scheduler
+        """
+
+    def __repr__(self) -> str:
+        """
+        Returns the string representation of the scheduler.
+
+        Returns
+        -------
+        str
+            String representation of the scheduler
+        """
+        return (f'{self.__class__.__name__}(' +
+                (f'{self.extra_repr()})' if self.extra_repr() else ')'))
+
+    @abstractmethod
+    def step(self, epoch: int | float) -> None:
+        """
+        Steps the scheduler.
+
+        Parameters
+        ----------
+        epoch : int | float
+            Current epoch
+        """
+
+    def extra_repr(self) -> str:
+        """
+        Additional string representation of the scheduler.
+
+        Returns
+        -------
+        str
+            Additional string representation of the scheduler
+        """
+        return ''
+
+    def get_step(self, epoch: int | float) -> float:
+        """
+        Returns the current value of the scheduler for the given epoch.
+
+        Parameters
+        ----------
+        epoch : int | float
+            Current epoch
+
+        Returns
+        -------
+        float
+            Current value of the scheduler for the given epoch
+        """
+        self.step(epoch)
+        return self()
+
+
+class BaseFuncScheduler(BaseScheduler, ABC):
+    """
+    Abstract base class for simple function schedulers.
+
+    Attributes
+    ----------
+    max_epochs : int
+        Maximum number of epochs
+    warmup : int
+        Number of warmup epochs
+    initial_val : float
+        Initial value of the scheduler
+    final_val : float
+        Final value of the scheduler
+    """
+    max_epochs: int
+    warmup: int
+    final_val: float
+    initial_val: float
+    _value: float
+
+    def __init__(
+            self,
+            max_epochs: int,
+            *,
+            warmup_epochs: int = 0,
+            final_value: float = 1,
+            initial_value: float = 0) -> None:
+        """
+        Parameters
+        ----------
+        max_epochs : int
+            Maximum number of epochs
+        warmup_epochs : int, Optional
+            Number of warmup epochs, default = 0
+        final_value : float, Optional
+            Maximum value of the scheduler, default = 1
+        initial_value : float, Optional
+            Minimum value of the scheduler, default = 0
+        """
+        super().__init__()
+        self.max_epochs = max_epochs
+        self.warmup = warmup_epochs
+        self.final_val = final_value
+        self.initial_val = initial_value
+        self._value = self.final_val
+        self.step(0)
+
+        if self.max_epochs < warmup_epochs:
+            raise ValueError(f'max_epochs ({self.max_epochs}) should be larger than warmup_epochs '
+                             f'({self.warmup})')
+
+    def __call__(self) -> float:
+        return self._value
+
+    def __getstate__(self) -> dict[str, Any]:
+        return {
+            'max_epochs': self.max_epochs,
+            'warmup': self.warmup,
+            'final_val': self.final_val,
+            'initial_val': self.initial_val,
+        }
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._value = self.final_val
+        self.step(0)
+
+    def extra_repr(self) -> str:
+        return (f'max_epochs={self.max_epochs}, warmup_epochs={self.warmup}, '
+                f'final_value={self.final_val}, initial_value={self.initial_val}')
+
+class LinearScheduler(BaseFuncScheduler):
+    """
+    Linear scheduler.
+
+    Attributes
+    ----------
+    max_epochs : int
+        Maximum number of epochs
+    warmup : int
+        Number of warmup epochs
+    initial_val : float
+        Initial value of the scheduler
+    final_val : float
+        Final value of the scheduler
+    """
+    def step(self, epoch: int | float) -> None:
+        progress: float
+
+        if epoch >= self.max_epochs:
+            self._value = self.final_val
+        elif epoch <= self.warmup:
+            self._value = self.initial_val
+        else:
+            progress = (epoch - self.warmup) / (self.max_epochs - self.warmup)
+            self._value = progress * (self.final_val - self.initial_val) + self.initial_val
+
+
+class SigmoidScheduler(BaseFuncScheduler):
+    """
+    Sigmoid scheduler.
+
+    Attributes
+    ----------
+    max_epochs : int
+        Maximum number of epochs
+    warmup : int
+        Number of warmup epochs
+    initial_val : float
+        Initial value of the scheduler
+    final_val : float
+        Final value of the scheduler
+    compression : float
+        Compression factor for the sigmoid function, higher values result in a steeper curve
+    """
+    compression: float = 10
+
+    def __init__(
+            self,
+            max_epochs: int,
+            *,
+            warmup_epochs: int = 0,
+            final_value: float = 1,
+            initial_value: float = 0,
+            compression: float = 0) -> None:
+        """
+        Parameters
+        ----------
+        max_epochs : int
+            Maximum number of epochs
+        warmup_epochs : int, Optional
+            Number of warmup epochs, default = 0
+        final_value : float, Optional
+            Maximum value of the scheduler, default = 1
+        initial_value : float, Optional
+            Minimum value of the scheduler, default = 0
+        compression : float, Optional
+            Compression factor for the tanh function, higher values result in a steeper curve,
+            default = 10
+        """
+        super().__init__(
+            max_epochs,
+            warmup_epochs=warmup_epochs,
+            final_value=final_value,
+            initial_value=initial_value,
+        )
+        self.compression = compression or self.compression
+
+    def __getstate__(self) -> dict[str, Any]:
+        return super().__getstate__() | {'compression': self.compression}
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        super().__setstate__(state)
+        self.compression = state['compression']
+
+    def extra_repr(self) -> str:
+        return f'{super().extra_repr()}, compression={self.compression}'
+
+    def step(self, epoch: int | float) -> None:
+        progress: float
+
+        if epoch >= self.max_epochs:
+            self._value = self.final_val
+        elif epoch <= self.warmup:
+            self._value = self.initial_val
+        else:
+            progress = 1 / (1 + np.exp(
+                -self.compression * ((epoch - self.warmup) / (self.max_epochs - self.warmup) - 0.5)
+            ))
+            self._value = progress * (self.final_val - self.initial_val) + self.initial_val
+
+
+class StepScheduler(BaseScheduler):
+    """
+    Step scheduler.
+
+    Attributes
+    ----------
+    warmup : int
+        Number of warmup epochs
+    max_steps : int
+        Maximum number of steps
+    step_size : int
+        Number of epochs between each step
+    gamma : float
+        Multiplicative factor of the step
+    initial_val : float
+        Initial value of the scheduler
+    """
+    def __init__(
+            self,
+            warmup: int,
+            max_steps: int,
+            step_size: int,
+            gamma: float,
+            initial_value: float = 1) -> None:
+        """
+        Parameters
+        ----------
+        warmup : int
+            Number of warmup epochs
+        max_steps : int
+            Maximum number of steps
+        step_size : int
+            Number of epochs between each step
+        gamma : float
+            Multiplicative factor of the step
+        initial_value : float
+            Initial value of the scheduler, default = 1
+        """
+        super().__init__()
+        self.warmup: int = warmup
+        self.max_steps: int = max_steps
+        self.step_size: int = step_size
+        self.gamma: float = gamma
+        self.initial_val: float = initial_value
+        self._value: float = self.initial_val
+
+    def __call__(self) -> float:
+        return self._value
+
+    def __getstate__(self) -> dict[str, Any]:
+        return {
+            'warmup': self.warmup,
+            'max_steps': self.max_steps,
+            'step_size': self.step_size,
+            'gamma': self.gamma,
+            'initial_val': self.initial_val,
+        }
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._value = self.initial_val
+
+    def extra_repr(self) -> str:
+        return (f'warmup={self.warmup}, max_steps={self.max_steps}, step_size={self.step_size}, '
+                f'initial_value={self.initial_val}')
+
+    def step(self, epoch: int | float) -> None:
+        steps: int = min(max(int((epoch - self.warmup) // self.step_size), 0), self.max_steps)
+
+        if not steps:
+            self._value = self.initial_val
+            return
+
+        self._value = self.initial_val * (self.gamma ** steps)
+
+
+class TanhScheduler(SigmoidScheduler):
+    """
+    Tanh scheduler.
+
+    Attributes
+    ----------
+    max_epochs : int
+        Maximum number of epochs
+    warmup : int
+        Number of warmup epochs
+    initial_val : float
+        Initial value of the scheduler
+    final_val : float
+        Final value of the scheduler
+    compression : float
+        Compression factor for the tanh function, higher values result in a steeper curve
+    """
+    compression: float = 5
+
+    def __init__(
+            self,
+            max_epochs: int,
+            *,
+            warmup_epochs: int = 0,
+            final_value: float = 1,
+            initial_value: float = 0,
+            compression: float = 0) -> None:
+        """
+        Parameters
+        ----------
+        max_epochs : int
+            Maximum number of epochs
+        warmup_epochs : int, Optional
+            Number of warmup epochs, default = 0
+        final_value : float, Optional
+            Maximum value of the scheduler, default = 1
+        initial_value : float, Optional
+            Minimum value of the scheduler, default = 0
+        compression : float, Optional
+            Compression factor for the sigmoid function, higher values result in a steeper curve,
+            default = 5
+        """
+        super().__init__(
+            max_epochs,
+            warmup_epochs=warmup_epochs,
+            final_value=final_value,
+            initial_value=initial_value,
+            compression=compression,
+        )
+
+    def step(self, epoch: int | float) -> None:
+        progress: float
+
+        if epoch >= self.max_epochs:
+            self._value = self.final_val
+        elif epoch <= self.warmup:
+            self._value = self.initial_val
+        else:
+            progress = 2 / (1 + np.exp(
+                -self.compression * (epoch - self.warmup) / (self.max_epochs - self.warmup)
+            )) - 1
+            self._value = progress * (self.final_val - self.initial_val) + self.initial_val