netloader 3.11.2__tar.gz → 3.11.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: netloader
-Version: 3.11.2
+Version: 3.11.4
 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.11.4}/netloader/__init__.py

@@ -6,7 +6,7 @@ import logging
 import warnings
 
 
-__version__ = '3.11.2'
+__version__ = '3.11.4'
 __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.11.4}/netloader/architectures/base.py

@@ -62,6 +62,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
     """

{netloader-3.11.2 → netloader-3.11.4}/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.11.4}/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.11.2 → netloader-3.11.4}/netloader/models/misc.py

@@ -12,7 +12,7 @@ from netloader.utils import Config, TypedModuleList
 from netloader.network import BaseNetwork, CompatibleNetwork, Network
 
 
-class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
+class MultiNetwork(BaseNetwork[TypedModuleList[BaseNetwork]]):
     """
     Network class to hold multiple networks and train them together.
 
@@ -27,7 +27,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
         is the batch size
     kl_loss : Tensor
         KL divergence loss on the latent space of shape (1) and type float, if using a sample layer
-    layers : TypedModuleList[CompatibleNetwork | Network]
+    layers : TypedModuleList[BaseNetwork]
         Network construction
     """
     _cache: bool
@@ -36,14 +36,14 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
     def __init__(
             self,
             name: str,
-            *nets: nn.Module | nn.ModuleList | CompatibleNetwork | Network,
+            *nets: nn.Module | nn.ModuleList | BaseNetwork,
             save_outputs: bool = ...) -> None: ...
 
     @overload
     def __init__(
             self,
             name: str,
-            *nets: str | nn.Module | nn.ModuleList | CompatibleNetwork | Network | Config,
+            *nets: str | nn.Module | nn.ModuleList | BaseNetwork | Config,
             save_outputs: bool = ...,
             suppress_warning: bool = ...,
             root: str = ...,
@@ -55,7 +55,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
     def __init__(
             self,
             name: str,
-            *nets: str | nn.Module | nn.ModuleList | CompatibleNetwork | Network | Config,
+            *nets: str | nn.Module | nn.ModuleList | BaseNetwork | Config,
             save_outputs: bool = False,
             suppress_warning: bool = False,
             root: str = '',
@@ -68,7 +68,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
         ----------
         name : str
             Name of the network, used for saving
-        *nets: str | nn.Module | nn.ModuleList | CompatibleNetwork | Network | Config
+        *nets: str | nn.Module | nn.ModuleList | BaseNetwork | Config
             Networks to combine sequentially
         save_outputs : bool, Optional
             If outputs from each network should be saved, default = False
@@ -91,7 +91,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
         self.name = name
         self.layers = TypedModuleList()
         calc_in_shape: bool = True
-        net: str | nn.Module | nn.ModuleList | CompatibleNetwork | Network | Config
+        net: str | nn.Module | nn.ModuleList | BaseNetwork | Config
         builds: list[bool] = [isinstance(net, (str, Config)) for net in nets]
         out_shape: list[int] | None
         default: dict[str, Any] | None
@@ -150,15 +150,15 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
                 in_shape = out_shape
 
     @overload
-    def __getitem__(self, item: int) -> CompatibleNetwork | Network: ...
+    def __getitem__(self, item: int) -> BaseNetwork: ...
 
     @overload
-    def __getitem__(self, item: slice) -> TypedModuleList[CompatibleNetwork | Network]: ...
+    def __getitem__(self, item: slice) -> TypedModuleList[BaseNetwork]: ...
 
     def __getitem__(
             self,
             item: int | slice,
-    ) -> CompatibleNetwork | Network | TypedModuleList[CompatibleNetwork | Network]:
+    ) -> BaseNetwork | TypedModuleList[BaseNetwork]:
         """
         Returns the network(s) at the given index or slice.
 
@@ -169,7 +169,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
 
         Returns
         -------
-        CompatibleNetwork | Network | TypedModuleList[CompatibleNetwork | Network]
+        BaseNetwork | TypedModuleList[BaseNetwork]
             The network(s) at the given index or slice
         """
         return self.layers[item]
@@ -183,7 +183,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
     def __setstate__(self, state: dict[str, Any]) -> None:
         super().__setstate__(state)
         self._cache = state.get('save_outputs', False)
-        self.layers = TypedModuleList(state.get('layers', []))
+        self.layers = TypedModuleList(state.get('layers', state.get('net')))
 
     def forward(self, x: TensorListLike) -> TensorListLike:
         """
@@ -200,7 +200,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
         TensorListLike
             Output tensor from the network with shape (N,...) and type float
         """
-        net: CompatibleNetwork | Network
+        net: BaseNetwork
         self.checkpoints = []
 
         for i, net in enumerate(self.layers):
@@ -231,11 +231,14 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
         dict[str, Any] | Config
             Network configuration
         """
-        net: CompatibleNetwork | Network
+        net: BaseNetwork
         config = Config(layers=[net.get_config() for net in self.layers])
         return config.to_dict() if dict_ else config
 
     def to(self, *args: Any, **kwargs: Any) -> Self:
         super().to(*args, **kwargs)
-        self.layers.to(*args, **kwargs)
+        layer: BaseNetwork
+
+        for layer in self.layers:
+            layer.to(*args, **kwargs)
         return self

{netloader-3.11.2 → netloader-3.11.4}/netloader/network.py

@@ -81,7 +81,7 @@ class BaseNetwork(nn.Module, ABC, Generic[ModuleListT]):
         """
         if item == 'net':
             warn(
-                'CompatibleNetwork.net is deprecated, please use CompatibleNetwork.layers '
+                'BaseNetwork.net is deprecated, please use BaseNetwork.layers '
                 'instead',
                 DeprecationWarning,
                 stacklevel=2,
@@ -100,6 +100,28 @@ class BaseNetwork(nn.Module, ABC, Generic[ModuleListT]):
         """
         return {'name': self.name, 'version': netloader.__version__}
 
+    def __setattr__(self, key: str, value: Any) -> None:
+        """
+        Sets the attribute of the network.
+
+        Parameters
+        ----------
+        key : str
+            Name of the attribute to set
+        value : Any
+            Value to set the attribute to
+        """
+        if key == 'net':
+            warn(
+                'BaseNetwork.net is deprecated, please use BaseNetwork.layers '
+                'instead',
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            super().__setattr__('layers', value)
+        else:
+            super().__setattr__(key, value)
+
     def __setstate__(self, state: dict[str, Any]) -> None:
         """
         Sets the state of the network for pickling
@@ -181,7 +203,14 @@ class BaseNetwork(nn.Module, ABC, Generic[ModuleListT]):
 
     def to(self, *args: Any, **kwargs: Any) -> Self:
         super().to(*args, **kwargs)
-        self.kl_loss = self.kl_loss.to(*args, **kwargs)
+        self.kl_loss = self.kl_loss.detach().to(*args, **kwargs)
+
+        for i, checkpoint in enumerate(self.checkpoints):
+            if hasattr(checkpoint, 'to'):
+                self.checkpoints[i] = checkpoint.detach().to(*args, **kwargs)
+            else:
+                self._logger.warning(f'Failed to move checkpoint {checkpoint.__class__.__name__} '
+                                     f'to device')
         return self
 
 
@@ -522,13 +551,6 @@ class Network(BaseNetwork[TypedModuleList[layers.BaseLayer]]):
 
         for layer in self.layers:
             layer.to(*args, **kwargs)
-
-        for i, checkpoint in enumerate(self.checkpoints):
-            if hasattr(checkpoint, 'to'):
-                self.checkpoints[i] = checkpoint.to(*args, **kwargs)
-            else:
-                self._logger.warning(f'Failed to move checkpoint {checkpoint.__class__.__name__} '
-                                     f'to device')
         return self
 
 

{netloader-3.11.2 → netloader-3.11.4}/netloader.egg-info/PKG-INFO

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