"""
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
[docs]
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,
}
[docs]
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,
}
[docs]
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}
[docs]
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
[docs]
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
[docs]
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
[docs]
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
[docs]
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']