netloader 3.11.0__tar.gz → 3.11.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: netloader
-Version: 3.11.0
+Version: 3.11.2
 Summary: Utility to generate and train PyTorch neural network objects from JSON files
 Author-email: Ethan Tregidga <ethan.tregidga@epfl.ch>
 License-Expression: MIT
@@ -9,14 +9,14 @@ Project-URL: Issues, https://github.com/EthanTreg/PyTorch-Network-Loader/issues
 Project-URL: ReadTheDocs, https://pytorch-network-loader.readthedocs.io/en/latest/
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.12
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
-Requires-Dist: numpy>=2.4.0
-Requires-Dist: torch>=2.11.0
+Requires-Dist: numpy>=2.0.0
+Requires-Dist: torch>=2.8.0
 Requires-Dist: packaging>=24.0
 Provides-Extra: flows
-Requires-Dist: zuko; extra == "flows"
+Requires-Dist: zuko>=1.5.0; extra == "flows"
 Provides-Extra: wandb
 Requires-Dist: wandb; extra == "wandb"
 Dynamic: license-file
@@ -37,16 +37,14 @@ For a real-world example of this package, see [Bayesian-DARKSKIES](https://githu
 
 ### Using Within Projects
 
-- pip install `netloader @ git+https://github.com/EthanTreg/PyTorch-Network-Loader@LATEST-VERSION`[^1] to
-  `requirements.txt`
-- Install using `pip install -r requirements.txt`
+- Install using `pip install netloader`[^1]
 - Example of [InceptionV4](https://arxiv.org/abs/1602.07261) can be downloaded under
   `./network_configs/inceptionv4.json` along with the composite layers in
   `./network_configs/composite_layers/`
 
 [^1]: To use normalising flows or weights & biases, `netloader` must be pip installed with the optional argument
 `flows` and/or `wandb`:
-`pip install netloader[flows] @ git+https://github.com/EthanTreg/PyTorch-Network-Loader@LATEST-VERSION`  
+`pip install netloader[flows,wandb]`
 
 ### Locally Running NetLoader
 

{netloader-3.11.0 → netloader-3.11.2}/README.md

@@ -14,16 +14,14 @@ For a real-world example of this package, see [Bayesian-DARKSKIES](https://githu
 
 ### Using Within Projects
 
-- pip install `netloader @ git+https://github.com/EthanTreg/PyTorch-Network-Loader@LATEST-VERSION`[^1] to
-  `requirements.txt`
-- Install using `pip install -r requirements.txt`
+- Install using `pip install netloader`[^1]
 - Example of [InceptionV4](https://arxiv.org/abs/1602.07261) can be downloaded under
   `./network_configs/inceptionv4.json` along with the composite layers in
   `./network_configs/composite_layers/`
 
 [^1]: To use normalising flows or weights & biases, `netloader` must be pip installed with the optional argument
 `flows` and/or `wandb`:
-`pip install netloader[flows] @ git+https://github.com/EthanTreg/PyTorch-Network-Loader@LATEST-VERSION`  
+`pip install netloader[flows,wandb]`
 
 ### Locally Running NetLoader
 

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

@@ -6,7 +6,7 @@ import logging
 import warnings
 
 
-__version__ = '3.11.0'
+__version__ = '3.11.2'
 __author__ = 'Ethan Tregidga'
 logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.WARNING)
 warnings.filterwarnings('once', category=DeprecationWarning, module=r'^netloader(\.|$)')
@@ -43,8 +43,23 @@ try:
 
     # Adds PyTorch Network Loader classes to list of safe PyTorch classes when loading saved
     # architectures
-    safe_globals(__name__, [models, architectures, loss_funcs, schedulers, transforms])
-    torch.serialization.add_safe_globals([network.Network, network.CompatibleNetwork])
+    safe_globals(
+        __name__,
+        [models, loss_funcs, schedulers, transforms],
+        use_all=True,
+    )
+    safe_globals(
+        __name__,
+        [architectures],
+        use_all=True,
+        legacy_path_replace=('architectures', 'networks'),
+    )
+    torch.serialization.add_safe_globals([
+        network.BaseNetwork,
+        network.Network,
+        network.CompatibleNetwork,
+    ])
+    torch.serialization.add_safe_globals([])
 
     __all__ = [
         'utils',

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

@@ -6,7 +6,8 @@ import logging as log
 from time import time
 from warnings import warn
 from itertools import repeat
-from typing import TYPE_CHECKING, Any, Self, Generic, Literal, Callable, cast, overload
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Self, Sequence, Generic, Literal, Callable, cast, overload
 
 import torch
 import numpy as np
@@ -16,17 +17,12 @@ from torch.utils.data import DataLoader
 from torch.optim.optimizer import ParamsT
 from torch._dynamo import OptimizedModule
 
-if TYPE_CHECKING:
-    from wandb import Run
-else:
-    Run = Any
-
 import netloader
 from netloader import utils
 from netloader.transforms import BaseTransform
 from netloader.architectures.utils import UtilityMixin
-from netloader.network import Network, CompatibleNetwork
-from netloader.data import Data, DataList, data_collation
+from netloader.data import ApplyFunc, Data, DataList, data_collation
+from netloader.network import BaseNetwork, CompatibleNetwork
 from netloader.utils.types import (
     TensorLike,
     NDArrayLike,
@@ -37,8 +33,13 @@ from netloader.utils.types import (
     TensorLossCT,
 )
 
+if TYPE_CHECKING:
+    from wandb import Run
+else:
+    Run = Any
+
 
-class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
+class BaseArchitecture(UtilityMixin, ABC, Generic[LossCT, TensorLossCT]):
     # pylint: disable=line-too-long
     """
     Base architecture class that other types of architectures build from
@@ -61,15 +62,38 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
-    net : Network | CompatibleNetwork
+    net : BaseNetwork
         Neural network
     """
     # pylint: enable=line-too-long
+    _half: bool
+    _train_state: bool = True
+    _plot_active: bool = False
+    _save_freq: int
+    _epoch: int = 0
+    _save_path: str = ''
+    _verbose: Literal['epoch', 'full', 'plot', 'progress', None]
+    _loader_states: tuple[Tensor, Tensor] | None = None
+    _loss_weights: dict[str, float | Callable[[], float]]
+    _optimiser_kwargs: dict[str, Any]
+    _scheduler_kwargs: dict[str, Any]
+    _logger: log.Logger = log.getLogger(__name__)
+    _device: torch.device = torch.device('cpu')
+    description: str
+    version: str = netloader.__version__
+    losses: tuple[list[LossCT], list[LossCT]]
+    transforms: dict[str, Sequence[BaseTransform] | BaseTransform | None]
+    idxs: ndarray | None = None
+    optimiser: optim.Optimizer
+    scheduler: optim.lr_scheduler.LRScheduler
+    net: BaseNetwork
+    run: Run | None = None
+
     def __init__(
             self,
             save_num: int | str,
             states_dir: str,
-            net: nn.Module | Network | CompatibleNetwork,
+            net: nn.Module | BaseNetwork,
             *,
             overwrite: bool = False,
             mix_precision: bool = False,
@@ -88,7 +112,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
             File number or name to save the architecture
         states_dir : str
             Directory to save the architecture
-        net : Module | Network | CompatibleNetwork
+        net : Module | BaseNetwork
             Network to predict low-dimensional data
         overwrite : bool, Optional
             If saving can overwrite an existing save file, if False and file with the same name
@@ -97,7 +121,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
             If mixed precision should be used, default = False
         save_freq : int, Optional
             Frequency of epochs to save the architecture, default = 1
-        learning_rate : float, Optional
+        learning_rate : float | tuple[float, ...], Optional
             Optimiser initial learning rate, default = 1e-3
         description : str, Optional
             Description of the architecture
@@ -114,34 +138,21 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
         scheduler_kwargs : dict[str, Any] | None, Optional
             Optional keyword arguments to pass to init_scheduler
         """
-        self._train_state: bool = True
-        self._plot_active: bool = False
-        self._half: bool = mix_precision
-        self._epoch: int = 0
-        self._save_freq: int = save_freq
-        self._save_path: str = ''
-        self._verbose: Literal['epoch', 'full', 'plot', 'progress', None] = verbose
-        self._loader_states: tuple[Tensor, Tensor] | None = None
-        self._loss_weights: dict[str, float | Callable[[], float]] = {}
-        self._optimiser_kwargs: dict[str, Any] = optimiser_kwargs or {}
-        self._scheduler_kwargs: dict[str, Any] = scheduler_kwargs or {}
-        self._logger: log.Logger = log.getLogger(__name__)
-        self._device: torch.device = torch.device('cpu')
-        self.description: str = description
-        self.version: str = netloader.__version__
-        self.losses: tuple[list[LossCT], list[LossCT]] = ([], [])
-        self.transforms: dict[str, list[BaseTransform] | BaseTransform | None] = {
+        self._half = mix_precision
+        self._save_freq = save_freq
+        self._verbose = verbose
+        self._loss_weights = {}
+        self._optimiser_kwargs = optimiser_kwargs or {}
+        self._scheduler_kwargs = scheduler_kwargs or {}
+        self.description = description
+        self.losses = ([], [])
+        self.transforms = {
             'ids': None,
             'inputs': in_transform,
             'targets': transform,
             'preds': transform,
         }
-        self.idxs: ndarray | None = None
-        self.optimiser: optim.Optimizer
-        self.scheduler: optim.lr_scheduler.LRScheduler
-        self.net: Network | CompatibleNetwork = net \
-            if isinstance(net, (Network, CompatibleNetwork)) else CompatibleNetwork(net=net)
-        self.run: Run | None = None
+        self.net = net if isinstance(net, BaseNetwork) else CompatibleNetwork(net=net)
 
         if save_num:
             self._save_path = utils.save_name(save_num, states_dir, self.net.name)
@@ -159,7 +170,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
                 DeprecationWarning,
                 stacklevel=2,
             )
-            self.init_optimiser = self.set_optimiser
+            self.init_optimiser = self.set_optimiser  # type: ignore[method-assign]
 
         if self._method_override('set_scheduler', UtilityMixin):
             warn(
@@ -168,7 +179,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
                 DeprecationWarning,
                 stacklevel=2,
             )
-            self.init_scheduler = self.set_scheduler
+            self.init_scheduler = self.set_scheduler  # type: ignore[method-assign]
 
         self.optimiser = self.init_optimiser(
             self.get_param_groups(learning_rate),
@@ -177,10 +188,11 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
         self.scheduler = self.init_scheduler(self.optimiser, **self._scheduler_kwargs)
 
         if isinstance(self.scheduler, optim.lr_scheduler.ReduceLROnPlateau):
-            self._scheduler_kwargs = ({
-                                          'factor': 0.5,
-                                          'min_lr': (learning_rate if isinstance(learning_rate, float) else learning_rate[0]) * 1e-3,
-                                      } | self._scheduler_kwargs)
+            self._scheduler_kwargs = {
+                'factor': 0.5,
+                'min_lr': (learning_rate if isinstance(learning_rate, float) else learning_rate[0])
+                          * 1e-3
+            } | self._scheduler_kwargs
             self.scheduler.load_state_dict(self._scheduler_kwargs)
 
         # Adds all architecture classes to list of safe PyTorch classes when loading saved
@@ -296,7 +308,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
                 DeprecationWarning,
                 stacklevel=2,
             )
-            self.init_optimiser = self.set_optimiser
+            self.init_optimiser = self.set_optimiser  # type: ignore[method-assign]
 
         if self._method_override('set_scheduler', UtilityMixin):
             warn(
@@ -305,7 +317,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
                 DeprecationWarning,
                 stacklevel=2,
             )
-            self.init_scheduler = self.set_scheduler
+            self.init_scheduler = self.set_scheduler  # type: ignore[method-assign]
 
         if isinstance(state['optimiser'], dict):
             self.optimiser = self.init_optimiser(
@@ -494,7 +506,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
                         DeprecationWarning,
                         stacklevel=2,
                     )
-                    loss = self._loss_tensor(in_data, target)
+                    loss = self._loss_tensor(in_data, target)  # type: ignore[call-arg]
 
         if isinstance(loss, dict) and 'total' not in loss:
             loss['total'] = self._loss_total(loss)
@@ -523,6 +535,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
         """
         raise DeprecationWarning
 
+    @abstractmethod
     def _loss_tensor(
             self,
             in_data: TensorListLike,
@@ -545,7 +558,6 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
         TensorLossCT
             Loss of shape (1) and type float or dictionary of losses of shape (1) and type float
         """
-        raise NotImplementedError
 
     def _loss_total(self, losses: dict[str, Tensor]) -> Tensor:
         """
@@ -666,7 +678,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
                         DeprecationWarning,
                         stacklevel=2,
                     )
-                    batch_loss = self._loss(in_data, target)
+                    batch_loss = self._loss(in_data, target)  # type: ignore[call-arg] # pylint: disable=no-value-for-parameter
 
                 if isinstance(batch_loss, dict) and loss:
                     for key, value in batch_loss.items():
@@ -680,27 +692,14 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
                     loss = batch_loss
 
                 if self._train_state:
-                    try:
-                        self._step(
-                            True,
-                            self._epoch + (i + 1) / len(loader),
-                            batch_loss['total'] if isinstance(batch_loss, dict) else batch_loss,
-                            dataset=loader.dataset.dataset
-                            if hasattr(loader, 'dataset') and hasattr(loader.dataset, 'dataset')
-                            else None,
-                        )
-                    except TypeError:
-                        warn(
-                            '_step without dataset parameter is deprecated, please update'
-                            'the method to include the dataset parameter',
-                            DeprecationWarning,
-                            stacklevel=2,
-                        )
-                        self._step(
-                            True,
-                            self._epoch + (i + 1) / len(loader),
-                            batch_loss['total'] if isinstance(batch_loss, dict) else batch_loss,
-                        )
+                    self._step(
+                        True,
+                        self._epoch + (i + 1) / len(loader),
+                        batch_loss['total'] if isinstance(batch_loss, dict) else batch_loss,
+                        dataset=loader.dataset.dataset
+                        if hasattr(loader, 'dataset') and hasattr(loader.dataset, 'dataset')
+                        else None,
+                    )
 
                 self._batch_print(i, time() - t_initial, loader, batch_loss)
 
@@ -746,8 +745,8 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
         metrics : float | None, Optional
             Loss metric to update ReduceLROnPlateau
         """
-        learning_rate: list[float]
-        new_learning_rate: list[float]
+        learning_rate: list[float | Tensor]
+        new_learning_rate: list[float | Tensor]
 
         if batch_step is None:
             warn(
@@ -812,9 +811,9 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
             Optional keyword arguments to pass to torch.compile
         """
         if level == 'net':
-            self.net = cast(Network | CompatibleNetwork, torch.compile(self.net, **kwargs))
+            self.net = cast(BaseNetwork, torch.compile(self.net, **kwargs))
         elif level == 'loss':
-            self._loss_tensor = torch.compile(self._loss_tensor, **kwargs)  # type: ignore[method-assign] # pylint: disable=attribute-defined-outside-init
+            self._loss_tensor = torch.compile(self._loss_tensor, **kwargs)  # type: ignore[method-assign] # pylint: disable=attribute-defined-outside-init, method-hidden
         else:
             raise ValueError(f'Invalid compile level ({level}), must be "net" or "loss"')
 
@@ -1023,11 +1022,13 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
         high_dim: list[Tensor] | list[Data[Tensor]] | list[DataList[Tensor | Data[Tensor]]]
         data: list[list[NDArrayLike | None]] = []
         data_: dict[str, NDArrayLike] = {}
-        transform: list[BaseTransform] | BaseTransform | None
-        transforms: dict[str, list[BaseTransform] | BaseTransform | None] = {
-            key: transform for key, transform in self.transforms.items()
-            if inputs or key != 'inputs'
-        }
+        transform: (Sequence[BaseTransform[ndarray, ndarray]] |
+                    BaseTransform[ndarray, ndarray] | None)
+        transforms: dict[
+            str,
+            Sequence[BaseTransform[ndarray, ndarray]] | BaseTransform[ndarray, ndarray] | None
+        ] = {key: transform for key, transform in self.transforms.items()
+             if inputs or key != 'inputs'}
         datum: NDArrayLike
         target: TensorLike
         in_data: TensorLike
@@ -1075,7 +1076,10 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
             datum = data_collation(list(datum_), data_field=True)
 
             if isinstance(datum, DataList) and transform:
-                data_[key] = datum.apply(transform, back=True).numpy()
+                data_[key] = datum.apply(cast(
+                    ApplyFunc[ndarray | Data[ndarray], ..., ndarray | Data[ndarray]],
+                    transform,
+                ), back=True).numpy()
                 data_[key] = DataList(
                     [trans(val, back=True) for val, trans in zip(
                         datum,
@@ -1198,7 +1202,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
         cuda_state: Tensor | None
         np.random.set_state(states.pop('numpy', np.random.get_state()))
         self._loader_states = states.pop('loaders', None)
-        torch.set_rng_state(states.pop('pytorch', torch.get_rng_state()))
+        torch.set_rng_state(states.pop('pytorch', torch.get_rng_state()).cpu())
 
         if torch.cuda.is_available() and (cuda_state := states.pop('cuda_rng', None)):
             torch.cuda.set_rng_state_all(cuda_state)
@@ -1222,7 +1226,7 @@ class BaseArchitecture(UtilityMixin, Generic[LossCT, TensorLossCT]):
         self._device, *_ = torch._C._nn._parse_to(*args, **kwargs)  # pylint: disable=protected-access
 
         if isinstance(self.net, OptimizedModule):
-            self.net._orig_mod.to(*args, **kwargs)
+            self.net._orig_mod.to(*args, **kwargs)  # pylint: disable=protected-access
         return self
 
     def train(self, train: bool) -> None:

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

@@ -11,8 +11,8 @@ from torch import Tensor, nn
 
 from netloader.data import DataList
 from netloader.utils import label_change
+from netloader.network import BaseNetwork
 from netloader.transforms import BaseTransform
-from netloader.network import Network, CompatibleNetwork
 from netloader.architectures.base import BaseArchitecture
 from netloader.loss_funcs import BaseLoss, MSELoss, CrossEntropyLoss
 from netloader.utils.types import TensorListLike, NDArrayListLike, TensorT
@@ -42,14 +42,14 @@ class BaseEncoder(BaseArchitecture):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
-    net : Network | CompatibleNetwork
+    net : BaseNetwork
         Neural network
     """
     def __init__(
             self,
             save_num: int | str,
             states_dir: str,
-            net: nn.Module | Network | CompatibleNetwork,
+            net: nn.Module | BaseNetwork,
             *,
             overwrite: bool = False,
             mix_precision: bool = False,
@@ -69,7 +69,7 @@ class BaseEncoder(BaseArchitecture):
             File number or name to save the architecture
         states_dir : str
             Directory to save the architecture
-        net : Module | Network | CompatibleNetwork
+        net : Module | BaseNetwork
             Network to predict low-dimensional data
         overwrite : bool, Optional
             If saving can overwrite an existing save file, if True and file with the same name
@@ -223,7 +223,7 @@ class Autoencoder(BaseArchitecture):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
-    net : Network | CompatibleNetwork
+    net : BaseNetwork
         Neural network
     reconstruct_func : BaseLoss
         Loss function for the reconstruction loss
@@ -234,7 +234,7 @@ class Autoencoder(BaseArchitecture):
             self,
             save_num: int | str,
             states_dir: str,
-            net: nn.Module | Network | CompatibleNetwork,
+            net: nn.Module | BaseNetwork,
             *,
             overwrite: bool = False,
             mix_precision: bool = False,
@@ -252,7 +252,7 @@ class Autoencoder(BaseArchitecture):
             File number or name to save the architecture
         states_dir : str
             Directory to save the architecture
-        net : Module | Network | CompatibleNetwork
+        net : Module | BaseNetwork
             Network to predict low-dimensional data
         overwrite : bool, Optional
             If saving can overwrite an existing save file, if True and file with the same name
@@ -441,14 +441,14 @@ class Decoder(BaseArchitecture):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
-    net : Network | CompatibleNetwork
+    net : BaseNetwork
         Neural network
     """
     def __init__(
             self,
             save_num: int | str,
             states_dir: str,
-            net: nn.Module | Network | CompatibleNetwork,
+            net: nn.Module | BaseNetwork,
             *,
             overwrite: bool = False,
             mix_precision: bool = False,
@@ -466,7 +466,7 @@ class Decoder(BaseArchitecture):
             File number or name to save the architecture
         states_dir : str
             Directory to save the architecture
-        net : Module | Network | CompatibleNetwork
+        net : Module | BaseNetwork
             Network to predict low-dimensional data
         overwrite : bool, Optional
             If saving can overwrite an existing save file, if True and file with the same name
@@ -601,7 +601,7 @@ class Encoder(BaseEncoder):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
-    net : Network | CompatibleNetwork
+    net : BaseNetwork
         Neural network
     """
     def batch_predict(self, data: TensorListLike, **_: Any) -> tuple[NDArrayListLike | None, ...]:

{netloader-3.11.0 → netloader-3.11.2}/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 Network, CompatibleNetwork
 from netloader.architectures.base import BaseArchitecture
 from netloader.architectures.encoder_decoder import BaseEncoder
 from netloader.utils.types import NDArrayLike, TensorListLike, NDArrayListLike, TensorT
@@ -29,7 +29,7 @@ class NormFlow(BaseArchitecture):
 
     Attributes
     ----------
-    net : Network | CompatibleNetwork
+    net : BaseNetwork
         Neural spline flow
     description : str
         Description of the architecture
@@ -146,14 +146,20 @@ class NormFlowEncoder(BaseEncoder):
         Architecture optimiser
     scheduler : LRScheduler
         Optimiser scheduler
-    net : Network | CompatibleNetwork
+    net : BaseNetwork
         Neural network
     """
+    _train_flow: bool
+    _train_encoder: bool
+    _checkpoint: int | None
+    _epochs: tuple[int, int]
+    net: BaseNetwork[nn.ModuleList]
+
     def __init__(
             self,
             save_num: int | str,
             states_dir: str,
-            net: nn.Module | Network | CompatibleNetwork,
+            net: nn.Module | BaseNetwork[nn.ModuleList],
             *,
             overwrite: bool = False,
             mix_precision: bool = False,
@@ -175,7 +181,7 @@ class NormFlowEncoder(BaseEncoder):
             File number or name to save the flow
         states_dir : str
             Directory to save the architecture and flow
-        net : Module | Network | CompatibleNetwork
+        net : Module | BaseNetwork[ModuleList]
             Normalising flow to predict low-dimensional data distribution
         overwrite : bool, Optional
             If saving can overwrite an existing save file, if True and file with the same name
@@ -226,11 +232,8 @@ class NormFlowEncoder(BaseEncoder):
             optimiser_kwargs=optimiser_kwargs,
             scheduler_kwargs=scheduler_kwargs,
         )
-        self._train_flow: bool
-        self._train_encoder: bool
-        self._checkpoint: int | None = net_checkpoint
-        self._epochs: tuple[int, int] = train_epochs
-
+        self._checkpoint = net_checkpoint
+        self._epochs = train_epochs
         self._train_flow = not self._epochs[0]
         self._train_encoder = bool(self._epochs[-1])
         self._loss_weights = {'flow': 1, 'encoder': 1}
@@ -240,6 +243,7 @@ class NormFlowEncoder(BaseEncoder):
             'max': transform,
             'meds': transform,
         }
+        assert isinstance(self.net, nn.ModuleList)
 
         if not self._train_encoder:
             self.net.layers[:-1].requires_grad_(False)
@@ -360,8 +364,8 @@ class NormFlowEncoder(BaseEncoder):
         Returns
         -------
         tuple[NDArrayListLike | None, ndarray | None]
-            Architecture output with shape (N,...) and type float and samples of shape (N,S) and type
-            float from each probability distribution
+            Architecture output with shape (N,...) and type float and samples of shape (N,S) and
+            type float from each probability distribution
         """
         samples: ndarray | None = None
         output: NormalizingFlow | TensorListLike | None = self.net(data)
@@ -411,6 +415,17 @@ class NormFlowEncoder(BaseEncoder):
             'train_epochs': self._epochs,
         }
 
+    def get_param_groups(self, learning_rate: float | tuple[float, ...] | None) -> ParamsT:
+        learning_rate = learning_rate or (0,) * 2
+
+        if isinstance(learning_rate, float):
+            learning_rate = (learning_rate,) * 2
+
+        return [
+            {'params': self.net.layers[:-1].parameters(), 'lr': learning_rate[0]},
+            {'params': self.net.layers[-1:].parameters(), 'lr': learning_rate[1]},
+        ]
+
     def predict(
             self,
             loader: DataLoader[Any],
@@ -475,27 +490,5 @@ class NormFlowEncoder(BaseEncoder):
         self._save_predictions(path, data)
         return data
 
-    def get_param_groups(self, learning_rate: tuple[float, float] | None) -> ParamsT:
-        """
-        Sets the optimiser for the architecture, by default AdamW.
-
-        Parameters
-        ----------
-        learning_rate : tuple[float, float] | None, Optional
-            Learning rate for the encoder and normalising flow, if None, learning rate is set to 0
-        **kwargs
-            Optional keyword arguments to pass to the optimiser
-
-        Returns
-        -------
-        Optimizer
-            Architecture optimiser
-        """
-        learning_rate = learning_rate or (0,) * 2
-        return [
-            {'params': self.net.layers[:-1].parameters(), 'lr': learning_rate[0]},
-            {'params': self.net.layers[-1:].parameters(), 'lr': learning_rate[1]},
-        ]
-
 
 __all__ = ['NormFlow', 'NormFlowEncoder']

{netloader-3.11.0 → netloader-3.11.2}/netloader/architectures/utils.py

@@ -98,7 +98,7 @@ class UtilityMixin:
         bool
             True if the method is overridden, False otherwise
         """
-        method: Callable = getattr(self, name, None)
+        method: Callable | None = getattr(self, name, None)
 
         if method is None:
             logger.warning(f'Method {name} not found in {type(self).__name__}')