netloader 3.11.0__py3-none-any.whl

netloader/layers/flows.py

@@ -0,0 +1,123 @@
+"""
+Normalising flow network layers
+"""
+from typing import Any
+
+from torch import Tensor
+from zuko.flows import NSF
+from zuko.distributions import NormalizingFlow
+
+from netloader.utils import Shapes
+from netloader.layers.base import BaseLayer
+
+
+class SplineFlow(BaseLayer):
+    """
+    Neural spline flow layer constructor.
+
+    Should only be used as the last layer in the network as it does not return a Tensor.
+
+    Attributes
+    ----------
+    group : int
+        Layer group, if 0 it will always be used, else it will only be used if its group matches the
+        Networks
+    description : str
+        Description of the layer
+    """
+    def __init__(
+            self,
+            transforms: int,
+            hidden_features: list[int],
+            net_out: list[int],
+            shapes: Shapes,
+            *,
+            context: bool = False,
+            features: int | None = None,
+            factor: float | None = None,
+            **kwargs: Any) -> None:
+        """
+        Generates a neural spline flow (NSF) for use in BaseNetwork
+
+        Adds attributes of name ('flow'), optimiser (Adam), and scheduler (ReduceLROnPlateau)
+
+        Parameters
+        ----------
+        transforms : int
+            Number of transforms
+        hidden_features : list[int]
+            Number of features in each of the hidden layers
+        net_out : list[int]
+            Shape of the network's output
+        shapes : Shapes
+            Shape of the outputs from each layer
+        context : bool, Optional
+            If the output from the previous layer should be used to condition the normalising flow,
+            default = False
+        features : int | None, Optional
+            Dimensions of the probability distribution, if factor is provided, features will not be
+            used
+        factor : float | None, Optional
+            Output features is equal to the factor of the network's output, will be used if
+            provided, else features will be used
+        **kwargs
+            Leftover parameters to pass to base layer for checking
+        """
+        super().__init__(**kwargs)
+        self._context: bool = context
+        self._transforms: int = transforms
+        self._hidden_features: list[int] = hidden_features
+        self._layer: NSF
+        context_: int
+        shape: list[int] = shapes[-1].copy()
+
+        # Number of features can be defined by either a factor of the output size or explicitly
+        if factor:
+            shape[-1] = max(1, int(net_out[-1] * factor))
+        elif features:
+            shape[-1] = features
+        else:
+            raise ValueError('Either features or factor must be provided and be non-zero')
+
+        if self._context:
+            context_ = shapes[-1][-1]
+        else:
+            context_ = 0
+
+        self._layer = NSF(
+            features=shape[-1],
+            context=context_,
+            transforms=transforms,
+            hidden_features=hidden_features,
+        )
+
+        shapes.append(shape)
+
+    def __getstate__(self) -> dict[str, Any]:
+        return super().__getstate__() | {
+            'context': self._context,
+            'transforms': self._transforms,
+            'features': self._layer.transform.transforms[0].passes,
+            'hidden_features': self._hidden_features,
+        }
+
+    def forward(self, x: Tensor, *_: Any, **__: Any) -> NormalizingFlow:
+        """
+        Forward pass of the neural spline flow layer
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor with shape (N,...) and type float, where N is the batch size
+
+        Returns
+        -------
+        NormalizingFlow
+            Normalising flow distribution
+        """
+        if self._context:
+            return self._layer(x)
+        return self._layer()
+
+
+__all__ = ['SplineFlow']

netloader/layers/linear.py

@@ -0,0 +1,445 @@
+"""
+Linear network layers
+"""
+from __future__ import annotations
+import logging as log
+from inspect import signature
+from typing import TYPE_CHECKING, Any, Self, Literal, cast
+
+import torch
+import numpy as np
+from torch import nn, Tensor
+
+from netloader.utils import Shapes, compare_versions
+from netloader.layers.base import BaseLayer, BaseSingleLayer
+
+if TYPE_CHECKING:
+    from netloader.network import Network
+
+
+class Activation(BaseSingleLayer):
+    """
+    Activation layer constructor.
+
+    Attributes
+    ----------
+    group : int
+        Layer group, if 0 it will always be used, else it will only be used if its group matches the
+        Networks
+    description : str
+        Description of the layer
+    layers : Sequential
+        Layers to loop through in the forward pass
+    """
+    def __init__(
+            self,
+            *,
+            activation: str = 'ELU',
+            shapes: Shapes | None = None,
+            activation_kwargs: dict[str, Any] | None = None,
+            **kwargs: Any) -> None:
+        """
+        Parameters
+        ----------
+        activation : str, Optional
+            Which activation function to use from PyTorch, default = 'ELU'
+        shapes : Shapes | None, Optional
+            Shape of the outputs from each layer, only required if tracking layer outputs is
+            necessary
+        activation_kwargs : dict[str, Any] | None, Optional
+            Additional keyword arguments to pass to the activation function
+        **kwargs
+            Leftover parameters to pass to base layer for checking
+        """
+        super().__init__(**({'idx': 0} | kwargs))
+        self._kwargs: dict[str, Any] = activation_kwargs or {}
+        self.layers.add_module('Activation', getattr(nn, activation)(
+            **({'inplace': True} if 'inplace' in signature(getattr(nn, activation)).parameters
+            else {}) | self._kwargs,
+        ))
+
+        # If not used as a layer in a network
+        if not shapes:
+            return
+
+        shapes.append(shapes[-1].copy())
+
+    def __getstate__(self) -> dict[str, Any]:
+        return super().__getstate__() | {
+            'activation': self.layers.Activation.__class__.__name__,
+            'activation_kwargs': self._kwargs,
+        }
+
+
+class Linear(BaseSingleLayer):
+    """
+    Linear layer constructor.
+
+    Attributes
+    ----------
+    group : int
+        Layer group, if 0 it will always be used, else it will only be used if its group matches the
+        Networks
+    description : str
+        Description of the layer
+    layers : Sequential
+        Layers to loop through in the forward pass
+    """
+    def __init__(
+            self,
+            net_out: list[int],
+            shapes: Shapes,
+            *,
+            features: int | None = None,
+            layer: int | None = None,
+            factor: float | None = None,
+            batch_norm: bool = False,
+            flatten_target: bool = False,
+            dropout: float = 0,
+            activation: str | None = 'SELU',
+            **kwargs: Any) -> None:
+        """
+        Parameters
+        ----------
+        net_out : list[int]
+            Shape of the network's output
+        shapes : Shapes
+            Shape of the outputs from each layer
+        features : int | None, Optional
+            Number of output features for the layer, if factor is provided, features will not be
+            used
+        layer : int | None, Optional
+            If factor is not None, which layer for factor to be relative to, if None, network output
+            will be used
+        factor : float | None, Optional
+            Output features is equal to the factor of the network's output, or if layer is provided,
+            which layer to be relative to, will be used if provided, else features will be used
+        batch_norm : bool, Optional
+            If batch normalisation should be used, default = False
+        flatten_target : bool, Optional
+            If the target should be flattened so that features is equal to the product of the target
+            multiplied by factor, if factor is provided, default = False
+        dropout : float, Optional
+            Probability of dropout, default =  0
+        activation : str | None, Optional
+            Which activation function to use from PyTorch, default = 'SELU'
+        **kwargs
+            Leftover parameters to pass to base layer for checking
+        """
+        super().__init__(**kwargs)
+        shape: list[int] = shapes[-1].copy()
+        target: list[int] = shapes[layer] if layer is not None else net_out
+
+        # Number of features can be defined by either a factor of the output size or explicitly
+        shape = self._check_factor_filters(
+            shape[::-1],
+            filters=features,
+            factor=factor,
+            target=[int(np.prod(target))] if flatten_target else target[::-1],
+        )[::-1]
+
+        self.layers.add_module('Linear', nn.Linear(
+            in_features=shapes[-1][-1],
+            out_features=shape[-1],
+        ))
+
+        # Optional layers
+        if activation:
+            self.layers.add_module('Activation', getattr(nn, activation)(
+                **{'inplace': True} if 'inplace' in signature(getattr(nn, activation)).parameters
+                else {},
+            ))
+
+        if batch_norm:
+            self.layers.add_module('BatchNorm', nn.BatchNorm1d(shape[0]))
+
+        if dropout:
+            self.layers.add_module('Dropout', nn.Dropout(dropout))
+
+        shapes.append(shape)
+
+    def __getstate__(self) -> dict[str, Any]:
+        return super().__getstate__() | {
+            'features': self.layers.Linear.out_features,  # type: ignore[union-attr]
+            'batch_norm': 'BatchNorm' in self.layers._modules,
+            'dropout': self.layers.Dropout.p if 'Dropout' in self.layers._modules else 0,  # type: ignore[union-attr]
+            'activation': self.layers.Activation.__class__.__name__
+                if 'Activation' in self.layers._modules else None,
+        }
+
+
+class OrderedBottleneck(BaseLayer):
+    """
+    Information-ordered bottleneck to randomly change the size of the bottleneck in an autoencoder
+    to encode the most important information in the first values of the latent space.
+
+    See `Information-Ordered Bottlenecks for Adaptive Semantic Compression
+    <https://arxiv.org/abs/2305.11213>`_ by Ho et al. (2023).
+
+    Attributes
+    ----------
+    group : int
+        Layer group, if 0 it will always be used, else it will only be used if its group matches the
+        Networks
+    min_size : int
+        Minimum gate size
+    description : str
+        Description of the layer
+    """
+    def __init__(self, shapes: Shapes, *, min_size: int = 0, **kwargs: Any) -> None:
+        """
+        Parameters
+        ----------
+        shapes : Shapes
+            Shape of the outputs from each layer
+        min_size : int, Optional
+            Minimum gate size, default = 0
+        **kwargs
+            Leftover parameters to pass to base layer for checking
+        """
+        super().__init__(**kwargs)
+        self._offset: bool = compare_versions(self._ver, '3.10.0')
+        self.min_size: int = min_size
+        shapes.append(shapes[-1].copy())
+
+    def __getstate__(self) -> dict[str, Any]:
+        return super().__getstate__() | {'min_size': self.min_size}
+
+    def forward(self, x: Tensor, *_: Any, **__: Any) -> Tensor:
+        """
+        Forward pass of the information-ordered bottleneck layer
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor with shape (N,...,Z) and type float, where N is the batch size and Z is the
+            latent space where Z > min_size
+
+        Returns
+        -------
+        Tensor
+            Input zeroed from a random index to the last value along the dimension Z with shape
+            (N,...,Z) and type float
+        """
+        if not self.training or self.min_size >= x.size(-1):
+            return x
+
+        idx: Tensor = torch.randint(
+            self.min_size,
+            x.size(-1) + int(self._offset),
+            (1,),
+            device=x.device,
+        )
+        gate: Tensor = torch.where(
+            torch.arange(x.size(-1), device=x.device) < idx + int(not self._offset),
+            torch.ones_like(idx),
+            torch.zeros_like(idx),
+        )
+        return x * gate
+
+    def extra_repr(self) -> str:
+        """
+        Displays layer parameters when printing the network
+
+        Returns
+        -------
+        str
+            Layer parameters
+        """
+        return f'min_size={self.min_size}'
+
+
+class Sample(BaseLayer):
+    """
+    Samples random values from a Gaussian distribution for a variational autoencoder,
+    mean and standard deviation are the first and second half of the input channels with the last
+    channel ignored if there are an odd number.
+
+    Attributes
+    ----------
+    group : int
+        Layer group, if 0 it will always be used, else it will only be used if its group matches the
+        Networks
+    description : str
+        Description of the layer
+    sample_layer : Normal
+        Layer to sample values from a Gaussian distribution
+    """
+    def __init__(self, idx: int, shapes: Shapes, **kwargs: Any) -> None:
+        """
+        Parameters
+        ----------
+        idx : int
+            Layer number
+        shapes : Shapes
+            Shape of the outputs from each layer
+        **kwargs
+            Leftover parameters to pass to base layer for checking
+        """
+        super().__init__(idx=idx, **kwargs)
+        self.sample_layer: torch.distributions.Normal = torch.distributions.Normal(
+            torch.tensor(0.).to(self._device),
+            torch.tensor(1.).to(self._device),
+        )
+
+        if shapes[-1][0] % 2 == 1:
+            log.warning(f'Sample in layer {idx} expects an even length along the first dimension, '
+                        f'but input shape is {shapes[-1]}, last element will be ignored')
+
+        shapes.append(shapes[-1].copy())
+        shapes[-1][0] = shapes[-1][0] // 2
+
+    def forward(self, x: Tensor, net: Network, *_: Any, **__: Any) -> Tensor:
+        """
+        Forward pass of the sampling layer for a variational autoencoder
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor with shape (N,C,...) | (N,Z) and type float, where N is the batch size, and
+            either the channels dimension, C, or latent, Z, containing the mean and standard
+            deviation
+        net : Network
+            Parent network that this layer is part of
+
+        Returns
+        -------
+        Tensor
+            Output tensor sampled from the input tensor split into mean and standard deviation with
+            shape (N,C/2,...) | (N,Z/2) and type float
+        """
+        split: int = x.size(1) // 2
+        mean: Tensor = x[:, :split]
+        std: Tensor = torch.exp(x[:, split:2 * split])
+        x = mean + std * self.sample_layer.rsample(mean.shape)
+
+        net.kl_loss = 0.5 * torch.mean(
+            mean ** 2 + std ** 2 - 2 * torch.log(std) - 1
+        )
+        return x
+
+    def to(self, *args: Any, **kwargs: Any) -> Self:
+        super().to(*args, **kwargs)
+        self.sample_layer.loc = self.sample_layer.loc.to(*args, **kwargs)
+        self.sample_layer.scale = self.sample_layer.scale.to(*args, **kwargs)
+        return self
+
+
+class Upsample(BaseSingleLayer):
+    """
+    Constructs an upsampler.
+
+    Attributes
+    ----------
+    group : int
+        Layer group, if 0 it will always be used, else it will only be used if its group matches the
+        Networks
+    description : str
+        Description of the layer
+    layers : Sequential
+        Layers to loop through in the forward pass
+    """
+    def __init__(
+            self,
+            idx: int,
+            shapes: Shapes,
+            *,
+            shape: list[int] | None = None,
+            scale: float | list[float] | tuple[float, ...] = 2,
+            mode: Literal['nearest', 'linear', 'bilinear', 'bicubic', 'trilinear'] = 'nearest',
+            **kwargs: Any) -> None:
+        """
+        Parameters
+        ----------
+        idx : int
+            Layer number
+        shapes : Shapes
+            Shape of the outputs from each layer
+        shape : list[int] | None, Optional
+            Shape of the output, will be used if provided, else scale will be used
+        scale : float | tuple[float] | tuple[float, ...], Optional
+            Factor to upscale all or individual dimensions, first dimension is ignored, won't be
+            used if shape is provided, default = 2
+        mode : {'nearest', 'linear', 'bilinear', 'bicubic', 'trilinear'}
+            What interpolation method to use for upsampling
+        **kwargs
+            Leftover parameters to pass to base layer for checking
+        """
+        super().__init__(idx=idx, **kwargs)
+        modes: dict[str, list[int]] = {
+            'nearest': [2, 3, 4],
+            'linear': [2],
+            'bilinear': [3],
+            'bicubic': [3],
+            'trilinear': [4],
+        }
+
+        # Check for errors
+        self._check_shape((2, 4), shapes[-1])
+        self._check_upsample(shapes[-1], shape)
+        self._check_options('mode', mode, set(modes))
+        self._check_mode_dimension(mode, shapes[-1], modes)
+
+        if isinstance(scale, list):
+            scale = tuple(scale)
+        elif not isinstance(scale, tuple):
+            scale = (scale,) * len(shapes[-1][1:])
+
+        if shape:
+            shapes.append(shape)
+        else:
+            shapes.append(shapes[-1].copy())
+            shapes[-1][1:] = [
+                int(length * factor) for length, factor in zip(shapes[-1][1:], scale)
+            ]
+
+        self.layers.add_module('Upsample', nn.Upsample(
+            size=tuple(shape) if shape else None,
+            scale_factor=None if shape else scale,
+            mode=mode,
+        ))
+
+    def __getstate__(self) -> dict[str, Any]:
+        layer: nn.Upsample = cast(nn.Upsample, self.layers.Upsample)
+        return super().__getstate__() | {
+            'shape': list(shape) if isinstance(shape := layer.size, tuple) else shape,
+            'scale': layer.scale_factor,
+            'mode': layer.mode,
+        }
+
+    @staticmethod
+    def _check_mode_dimension(mode: str, shape: list[int], modes: dict[str, list[int]]) -> None:
+        """
+        Checks if the upsampling mode supports the number of input dimensions
+
+        Parameters
+        ----------
+        mode : str
+            Current upsampling mode
+        shape : list[int]
+            Input shape
+        modes : dict[str, list[int]]
+            Modes with a list of supported dimensions
+        """
+        if len(shape) not in modes[mode]:
+            raise ValueError(f'{mode} only supports input dimensions of {modes[mode]}, '
+                             f'input shape is {shape}')
+
+    @staticmethod
+    def _check_upsample(in_shape: list[int], out_shape: list[int] | None) -> None:
+        """
+        Checks if the input shape has the same number of dimensions as the output shape
+
+        Parameters
+        ----------
+        in_shape : list[int]
+            Input shape
+        out_shape : list[int] | None
+            Target output shape
+        """
+        if out_shape is not None and len(out_shape) != 1 and len(out_shape) + 1 != len(in_shape):
+            raise ValueError(f'Target shape {out_shape} is not compatible with input shape '
+                             f'{in_shape}, check that channels dimension is not in target shape')
+
+
+__all__ = ['Activation', 'Linear', 'OrderedBottleneck', 'Sample', 'Upsample']