netloader 3.11.2__tar.gz → 3.12.0__tar.gz

{netloader-3.11.2 → netloader-3.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: netloader
-Version: 3.11.2
+Version: 3.12.0
 Summary: Utility to generate and train PyTorch neural network objects from JSON files
 Author-email: Ethan Tregidga <ethan.tregidga@epfl.ch>
 License-Expression: MIT

{netloader-3.11.2 → netloader-3.12.0}/netloader/__init__.py

@@ -6,7 +6,7 @@ import logging
 import warnings
 
 
-__version__ = '3.11.2'
+__version__ = '3.12.0'
 __author__ = 'Ethan Tregidga'
 logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.WARNING)
 warnings.filterwarnings('once', category=DeprecationWarning, module=r'^netloader(\.|$)')

{netloader-3.11.2 → netloader-3.12.0}/netloader/architectures/base.py

@@ -46,6 +46,8 @@ class BaseArchitecture(UtilityMixin, ABC, Generic[LossCT, TensorLossCT]):
 
     Attributes
     ----------
+    checkpoints : bool
+        If checkpoints should be saved after network forward pass
     description : str
         Description of the architecture
     version : str
@@ -62,6 +64,8 @@ class BaseArchitecture(UtilityMixin, ABC, Generic[LossCT, TensorLossCT]):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
+    run : Run | None
+        Weights & Biases run object for logging and tracking experiments
     net : BaseNetwork
         Neural network
     """
@@ -79,6 +83,7 @@ class BaseArchitecture(UtilityMixin, ABC, Generic[LossCT, TensorLossCT]):
     _scheduler_kwargs: dict[str, Any]
     _logger: log.Logger = log.getLogger(__name__)
     _device: torch.device = torch.device('cpu')
+    checkpoints: bool = False
     description: str
     version: str = netloader.__version__
     losses: tuple[list[LossCT], list[LossCT]]
@@ -283,7 +288,7 @@ class BaseArchitecture(UtilityMixin, ABC, Generic[LossCT, TensorLossCT]):
         self.net = state['net']
         self.run = state.get('run', None)
 
-        if not utils.compare_versions(self.version, netloader.__version__):
+        if not utils.compare_versions(self.version, netloader.__version__, level='minor'):
             warn(
                 f'Architecture version ({self.version}) is older than the current '
                 f'NetLoader version ({netloader.__version__}), please resave the architecture '
@@ -513,6 +518,9 @@ class BaseArchitecture(UtilityMixin, ABC, Generic[LossCT, TensorLossCT]):
 
         self._update(loss['total'] if isinstance(loss, dict) else loss)
 
+        if not self.checkpoints:
+            self.net.clear_checkpoints()
+
         if isinstance(loss, dict):
             return {key: value.item() for key, value in loss.items()}  # type: ignore[return-value]
         return loss.item()  # type: ignore[return-value]
@@ -1067,6 +1075,9 @@ class BaseArchitecture(UtilityMixin, ABC, Generic[LossCT, TensorLossCT]):
                 ]))
                 self._predict_print(i, time() - t_initial, loader)
 
+        if not self.checkpoints:
+            self.net.clear_checkpoints()
+
         # Transforms all data and saves it to a dictionary
         for (key, transform), datum_ in zip(transforms.items(), zip(*data)):
             if datum_[0] is None:
@@ -1085,6 +1096,7 @@ class BaseArchitecture(UtilityMixin, ABC, Generic[LossCT, TensorLossCT]):
                         datum,
                         transform if isinstance(transform, list) else repeat(transform),
                     )],
+                    names=datum.get_names(),
                 )
             elif isinstance(transform, BaseTransform):
                 assert not isinstance(datum, DataList)

{netloader-3.11.2 → netloader-3.12.0}/netloader/architectures/encoder_decoder.py

@@ -27,7 +27,7 @@ class BaseEncoder(BaseArchitecture):
     description : str
         Description of the architecture
     version : str
-        Version of the architecture
+        Version of the architecture when it was created or re-saved
     losses : tuple[list[LossCT], list[LossCT]]
         Architecture training and validation losses as a float or dictionary of losses for each loss
         function
@@ -42,6 +42,8 @@ class BaseEncoder(BaseArchitecture):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
+    run : Run | None
+        Weights & Biases run object for logging and tracking experiments
     net : BaseNetwork
         Neural network
     """
@@ -211,6 +213,8 @@ class Autoencoder(BaseArchitecture):
     ----------
     description : str
         Description of the architecture
+    version : str
+        Version of the architecture when it was created or re-saved
     losses : tuple[list[LossCT], list[LossCT]]
         Architecture training and validation losses as a float or dictionary of losses for each loss
         function
@@ -223,6 +227,8 @@ class Autoencoder(BaseArchitecture):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
+    run : Run | None
+        Weights & Biases run object for logging and tracking experiments
     net : BaseNetwork
         Neural network
     reconstruct_func : BaseLoss
@@ -427,6 +433,8 @@ class Decoder(BaseArchitecture):
     ----------
     description : str
         Description of the architecture
+    version : str
+        Version of the architecture when it was created or re-saved
     losses : tuple[list[LossCT], list[LossCT]]
         Architecture training and validation losses as a float or dictionary of losses for each loss
         function
@@ -441,6 +449,8 @@ class Decoder(BaseArchitecture):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
+    run : Run | None
+        Weights & Biases run object for logging and tracking experiments
     net : BaseNetwork
         Neural network
     """
@@ -587,6 +597,8 @@ class Encoder(BaseEncoder):
     ----------
     description : str
         Description of the architecture
+    version : str
+        Version of the architecture when it was created or re-saved
     losses : tuple[list[LossCT], list[LossCT]]
         Architecture training and validation losses as a float or dictionary of losses for each loss
         function
@@ -601,6 +613,8 @@ class Encoder(BaseEncoder):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
+    run : Run | None
+        Weights & Biases run object for logging and tracking experiments
     net : BaseNetwork
         Neural network
     """

{netloader-3.11.2 → netloader-3.12.0}/netloader/architectures/flows.py

@@ -14,8 +14,8 @@ from zuko.distributions import NormalizingFlow
 from netloader.data import Data
 from netloader.utils import label_change
 from netloader.loss_funcs import BaseLoss
-from netloader.network import BaseNetwork
 from netloader.transforms import BaseTransform
+from netloader.network import BaseNetwork, Network
 from netloader.architectures.base import BaseArchitecture
 from netloader.architectures.encoder_decoder import BaseEncoder
 from netloader.utils.types import NDArrayLike, TensorListLike, NDArrayListLike, TensorT
@@ -29,10 +29,10 @@ class NormFlow(BaseArchitecture):
 
     Attributes
     ----------
-    net : BaseNetwork
-        Neural spline flow
     description : str
         Description of the architecture
+    version : str
+        Version of the architecture when it was created or re-saved
     losses : tuple[list[LossCT], list[LossCT]]
         Architecture training and validation losses as a float or dictionary of losses for each loss
         function
@@ -45,6 +45,10 @@ class NormFlow(BaseArchitecture):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
+    run : Run | None
+        Weights & Biases run object for logging and tracking experiments
+    net : BaseNetwork
+        Neural spline flow
     """
     @staticmethod
     def _data_loader_translation(low_dim: TensorT, high_dim: TensorT) -> tuple[TensorT, TensorT]:
@@ -132,6 +136,8 @@ class NormFlowEncoder(BaseEncoder):
     ----------
     description : str
         Description of the architecture
+    version : str
+        Version of the architecture when it was created or re-saved
     losses : tuple[list[LossCT], list[LossCT]]
         Architecture training and validation losses as a float or dictionary of losses for each loss
         function
@@ -146,6 +152,8 @@ class NormFlowEncoder(BaseEncoder):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
+    run : Run | None
+        Weights & Biases run object for logging and tracking experiments
     net : BaseNetwork
         Neural network
     """
@@ -222,7 +230,7 @@ class NormFlowEncoder(BaseEncoder):
             net,
             overwrite=overwrite,
             mix_precision=mix_precision,
-            learning_rate=0,
+            learning_rate=0.,
             description=description,
             verbose=verbose,
             classes=classes,
@@ -243,7 +251,7 @@ class NormFlowEncoder(BaseEncoder):
             'max': transform,
             'meds': transform,
         }
-        assert isinstance(self.net, nn.ModuleList)
+        assert isinstance(self.net, Network)
 
         if not self._train_encoder:
             self.net.layers[:-1].requires_grad_(False)

netloader-3.12.0/netloader/data/__init__.py

@@ -0,0 +1,8 @@
+"""
+Base dataset classes and data structures for NetLoader.
+"""
+from netloader.data.datasets import BaseDataset, loader_init
+from netloader.data.structures import ApplyFunc, Data, DataList, data_collation
+
+
+__all__ = ['ApplyFunc', 'Data', 'DataList', 'BaseDataset', 'data_collation', 'loader_init']

netloader-3.12.0/netloader/data/datasets.py

@@ -0,0 +1,270 @@
+"""
+Base dataset classes for use with BaseArchitecture.
+"""
+from __future__ import annotations
+import logging as log
+from typing import Any, Generic, Literal, Sequence, cast, overload
+
+import torch
+import numpy as np
+from torch.utils.data import Dataset, DataLoader, Subset
+from numpy import ndarray
+
+from netloader.utils.types import ArrayLike, DatasetT, DataListT
+
+
+class BaseDatasetMeta(type):
+    """
+    Automatically creates an index for each sample in the dataset after the dataset has been
+    initialised.
+    """
+    def __call__(cls: type[DatasetT], *args: Any, **kwargs: Any) -> DatasetT:
+        """
+        Parameters
+        ----------
+        cls : type[DatasetT]
+            Class that inherited BaseDatasetMeta
+        *args
+            Optional arguments to pass to BaseDataset class
+        **kwargs
+            Optional keyword arguments to pass to BaseDataset class
+
+        Returns
+        -------
+        DatasetT
+            Dataset instance from the class that inherited BaseDatasetMeta
+        """
+        instance: DatasetT = type.__call__(cls, *args, **kwargs)
+
+        if hasattr(instance, 'idxs') and instance.idxs.dtype != np.int_:
+            raise ValueError(f'idxs attribute already exists and does not have type int '
+                             f'({instance.idxs.dtype}), idxs attribute must be reserved for sample '
+                             f'index')
+
+        if not hasattr(instance, 'high_dim') or instance.high_dim is None:
+            raise ValueError(f'{instance.__class__.__name__} has no high_dim attribute which is '
+                             f'required by BaseDatasetMeta for creating idxs attribute')
+
+        if len(instance.idxs) == 0:
+            instance.idxs = np.arange(len(instance.high_dim))
+        elif len(instance.idxs) != len(instance.high_dim):
+            log.getLogger(__name__).warning(f'Length of idxs ({len(instance.idxs)}) and length of '
+                                            f'high_dim ({len(instance.high_dim)}) does not match, '
+                                            f'idxs will be sent to a range of high_dim length')
+            instance.idxs = np.arange(len(instance.high_dim))
+
+        for attribute in ('extra', 'low_dim', 'high_dim'):
+            if (getattr(instance, attribute) is not None and
+                    len(getattr(instance, attribute)) != len(instance.idxs)):
+                raise ValueError(f'Length of attribute {attribute} '
+                                 f'({len(getattr(instance, attribute))}) and idxs '
+                                 f'({len(instance.idxs)}) does not match')
+        return instance
+
+
+class BaseDataset(
+    Dataset[tuple[int, DataListT, DataListT, Any]],
+    Generic[DataListT],
+    metaclass=BaseDatasetMeta):
+    """
+    Base dataset class for use with BaseNetwork.
+
+    Attributes
+    ----------
+    extra : list[Any] | ArrayLike | None
+        Additional data for each sample in the dataset of length N with shape (N,...) and type Any
+    idxs : ndarray
+        Index for each sample in the dataset with shape (N) and type int
+    low_dim : DataListT | None
+        Low dimensional data for each sample in the dataset with shape (N,...)
+    high_dim : DataListT | None
+        High dimensional data for each sample in the dataset with shape (N,...), this is required
+    """
+    def __init__(self) -> None:
+        super().__init__()
+        self.extra: list[Any] | ArrayLike | None = None
+        self.idxs: ndarray = np.array([], dtype=np.int_)
+        self.low_dim: DataListT | None = None
+        self.high_dim: DataListT | None = None
+
+    def __len__(self) -> int:
+        """
+        Returns the number of samples in the dataset
+
+        Returns
+        -------
+        int
+            Number of samples in the dataset
+        """
+        return len(self.idxs)
+
+    def __getitem__(self, idx: int) -> tuple[int, DataListT, DataListT, Any]:
+        """
+        Parameters
+        ----------
+        idx : int
+            Sample index
+
+        Returns
+        -------
+        tuple[int, DataListT, DataListT, Any]
+            Sample index, low dimensional data, high dimensional data, and extra data
+        """
+        return self.idxs[idx], self.get_low_dim(idx), self.get_high_dim(idx), self.get_extra(idx)
+
+    def get_extra(self, idx: int) -> Any:
+        """
+        Gets extra data for the sample of the given index
+
+        Parameters
+        ----------
+        idx : int
+            Sample index
+
+        Returns
+        -------
+        Any
+            Sample extra data
+        """
+        return torch.tensor(()) if self.extra is None else self.extra[idx]
+
+    def get_high_dim(self, idx: int) -> DataListT:
+        """
+        Gets a high dimensional sample of the given index
+
+        Parameters
+        ----------
+        idx : int
+            Sample index
+
+        Returns
+        -------
+        DataListT
+            High dimensional sample
+        """
+        assert self.high_dim is not None
+        return cast(DataListT, self.high_dim[idx])
+
+    def get_low_dim(self, idx: int) -> DataListT:
+        """
+        Gets a low dimensional sample of the given index
+
+        Parameters
+        ----------
+        idx : int
+            Sample index
+
+        Returns
+        -------
+        DataListT
+            Low dimensional sample
+        """
+        return cast(DataListT, torch.tensor(()) if self.low_dim is None else self.low_dim[idx])
+
+    def step(self, epoch: float) -> None:
+        """
+        Step the dataset for each training iteration.
+
+        Parameters
+        ----------
+        epoch : float
+            Current epoch number
+        """
+
+
+@overload
+def loader_init(
+        dataset: DatasetT,
+        *,
+        return_idxs: Literal[False],
+        batch_size: int = ...,
+        ratios: list[float] | tuple[float, ...] | None = ...,
+        idxs: list[ndarray] | tuple[ndarray, ...] | ndarray | None = ...,
+        **kwargs: Any) -> tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...]: ...
+
+@overload
+def loader_init(
+        dataset: DatasetT,
+        *,
+        return_idxs: Literal[True],
+        batch_size: int = ...,
+        ratios: list[float] | tuple[float, ...] | None = ...,
+        idxs: list[ndarray] | tuple[ndarray, ...] | ndarray | None = ...,
+        **kwargs: Any) -> tuple[
+    tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...],
+    tuple[Sequence[int], ...]]: ...
+
+def loader_init(
+        dataset: DatasetT,
+        *,
+        return_idxs: bool = False,
+        batch_size: int = 64,
+        ratios: list[float] | tuple[float, ...] | None = None,
+        idxs: list[ndarray] | tuple[ndarray, ...] | ndarray | None = None,
+        **kwargs: Any,
+) -> (tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...] |
+      tuple[
+          tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...],
+          tuple[Sequence[int], ...]]):
+    # pylint: disable=line-too-long
+    """
+    Initialises data loaders from a subset of the dataset with the given ratios.
+
+    Parameters
+    ----------
+    dataset : DatasetT
+        Dataset to create data loaders from
+    return_idxs : bool, Optional
+        If the indexes for each data loader should be returned, default = False
+    batch_size : int, Optional
+        Batch size when sampling from the data loaders, default = 64
+    ratios : list[float] | tuple[float, ...] | None, Optional
+        Ratios of length M to split up the dataset into subsets, if idxs is provided, dataset will
+        first be split up using idxs and ratios will be used on the remaining samples,
+        default = (0.8,0.2)
+    idxs : list[ndarray] | tuple[ndarray, ...] | ndarray | None, Optional
+        Dataset indexes for creating the subsets with shape (N,S), where N is the number of subsets
+        and S is the number of samples in each subset
+    **kwargs : Any
+        Optional keyword arguments to pass to DataLoader
+
+    Returns
+    -------
+    tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...] | tuple[tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...], tuple[Sequence[int], ...]]
+        Data loaders for each subset of length N + M and optionally indexes used for each subset
+    """
+    # pylint: enable=line-too-long
+    num: int
+    slice_: float | ndarray
+    loaders: list[DataLoader[tuple[int, DataListT, DataListT, Any]]] = []
+    data_idxs: ndarray = np.arange(len(dataset))
+    loader: DataLoader[tuple[int, DataListT, DataListT, Any]]
+    idxs = [] if idxs is None else list(idxs) \
+        if isinstance(idxs, (tuple, list)) or np.ndim(idxs) > 1 else [idxs]
+    ratios = (0.8, 1) if ratios is None else list(np.cumsum(np.array(ratios) / np.sum(ratios)))
+    np.random.shuffle(data_idxs)
+
+    for slice_ in idxs + list(ratios):
+        if isinstance(slice_, (int, float)):
+            num = max(int(len(data_idxs) * slice_), 1)
+            slice_ = data_idxs[:num]
+
+        if not np.isin(data_idxs, slice_).any():
+            continue
+
+        loaders.append(DataLoader(
+            Subset(
+                cast(Dataset[tuple[int, DataListT, DataListT, Any]], dataset),
+                data_idxs[np.isin(data_idxs, slice_)].tolist(),
+            ),
+            batch_size=batch_size,
+            **{'shuffle': True} | kwargs,
+        ))
+        data_idxs = np.delete(data_idxs, np.isin(data_idxs, slice_))
+
+    if return_idxs:
+        return tuple(loaders), tuple(cast(Subset, loader.dataset).indices for loader in loaders)
+    return tuple(loaders)
+
+
+__all__ = ['BaseDataset', 'loader_init']