"""
Convolutional network layers
"""
from inspect import signature
from typing import Any, Type, Literal
import torch
import numpy as np
from torch import nn, Tensor
from netloader.layers.misc import LayerNorm, Pad
from netloader.utils import Shapes, compare_versions
from netloader.layers.base import BaseLayer, BaseSingleLayer
from netloader.layers.utils import _int_list_conversion, _kernel_shape, _padding
[docs]
class Conv(BaseSingleLayer):
"""
Convolutional layer constructor.
Supports 1D, 2D, and 3D convolution.
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,
*,
filters: int | None = None,
layer: int | None = None,
factor: float | None = None,
groups: int = 1,
kernel: int | list[int] = 3,
stride: int | list[int] = 1,
padding: int | Literal['same'] | list[int] = 0,
dropout: float = 0,
activation: str | None = 'ELU',
norm: Literal[None, 'batch', 'layer'] = None,
padding_mode: Literal[None, 'replicate', 'zeros', 'reflect', 'circular'] = None,
**kwargs: Any) -> None:
"""
Parameters
----------
net_out : list[int]
Shape of the network's output, required only if layer contains factor
shapes : Shapes
Shape of the outputs from each layer
filters : int | None, Optional
Number of convolutional filters, will be used if provided, else factor will 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
Number of convolutional filters equal to the output channels, or if layer is provided,
the layer's channels, multiplied by factor, won't be used if filters is provided
groups : int, Optional
Number of input channel groups, each with its own convolutional filter(s), input and
output channels must both be divisible by the number of groups, default = 1
kernel : int | list[int], Optional
Size of the kernel, default = 3
stride : int | list[int], Optional
Stride of the kernel, default = 1
padding : int | {'same'} | list[int], Optional
Input padding, can an int, list of ints or 'same' where 'same' preserves the input
shape, default = 0
dropout : float, Optional
Probability of dropout, default = 0
activation : str | None, Optional
Which activation function to use from PyTorch, default = 'ELU'
norm : {None, 'batch', 'layer'}
If batch or layer normalisation should be used
padding_mode : {None, 'zeros', 'reflect', 'replicate', 'circular'}
Padding mode to use from PyTorch, if None zero padding is used if version is >3.9.4 else
replication padding is used
**kwargs
Leftover parameters to pass to base layer for checking
"""
super().__init__(**kwargs)
self._pad: int | Literal['same'] | list[int] = padding
asymmetric: bool = False
padding_: int | list[int]
shape: list[int] = shapes[-1].copy()
target: list[int] = shapes[layer] if layer is not None else net_out
conv: Type[nn.Module]
dropout_: Type[nn.Module]
batch_norm_: Type[nn.Module]
# Check for errors and calculate same padding
if isinstance(self._pad, str):
self._check_options('padding', self._pad, {'same'})
self._check_stride(stride)
asymmetric, padding_ = _padding(kernel, stride, shape, shape)
else:
padding_ = self._pad
if padding_mode is None:
padding_mode = 'zeros' if compare_versions(self._ver, '3.10.0') else \
'replicate'
if asymmetric and compare_versions(self._ver, '3.10.1'):
assert isinstance(padding_, list)
self.layers.add_module('Pad', Pad(
tuple(padding_),
mode='constant' if padding_mode == 'zeros' else padding_mode,
))
self._pad = 0
# Check for errors and calculate number of filters
self._check_shape((2, 4), shape)
self._check_factor_filters(shape, filters=filters, factor=factor, target=target)
self._check_groups(shapes[-1][0], shape[0], groups)
self._check_options('norm', norm, {None, 'batch', 'layer'})
conv, dropout_, batch_norm_ = [
[nn.Conv1d, nn.Dropout1d, nn.BatchNorm1d],
[nn.Conv2d, nn.Dropout2d, nn.BatchNorm2d],
[nn.Conv3d, nn.Dropout3d, nn.BatchNorm3d],
][len(shape) - 2]
self.layers.add_module('Conv', conv(
in_channels=shapes[-1][0],
out_channels=shape[0],
kernel_size=kernel,
stride=stride,
padding=self._pad,
groups=groups,
padding_mode=padding_mode,
))
# Optional layers
if activation:
self.layers.add_module('Activation', getattr(nn, activation)(
**{'inplace': True} if 'inplace' in signature(getattr(nn, activation)).parameters
else {},
))
if norm == 'batch':
self.layers.add_module('BatchNorm', batch_norm_(shape[0]))
elif norm == 'layer':
self.layers.add_module('LayerNorm', LayerNorm(shape=shape[0:1]))
if dropout:
self.layers.add_module('Dropout', dropout_(dropout))
if self._pad == 'same' or asymmetric:
shapes.append(shape)
else:
shapes.append(_kernel_shape(kernel, stride, padding_, shape))
def __getstate__(self) -> dict[str, Any]:
return super().__getstate__() | {
'filters': self.layers.Conv.out_channels, # type: ignore[union-attr]
'groups': self.layers.Conv.groups, # type: ignore[union-attr]
'kernel': self.layers.Conv.kernel_size, # type: ignore[union-attr]
'stride': self.layers.Conv.stride,
'padding': self._pad,
'dropout': self.layers.Dropout.p if hasattr(self.layers, 'Dropout') else 0, # type: ignore[union-attr]
'activation': self.layers.Activation.__class__.__name__
if hasattr(self.layers, 'Activation') else None,
'norm': 'batch' if hasattr(self.layers, 'BatchNorm') else
'layer' if hasattr(self.layers, 'LayerNorm') else None,
'padding_mode': self.layers.Conv.padding_mode, # type: ignore[union-attr]
}
@staticmethod
def _check_groups(in_channels: int, out_channels: int, groups: int) -> None:
"""
Checks if the number of groups is compatible with the number of input and output channels
Parameters
----------
in_channels : int
Number of input channels
out_channels : int
Number of output channels
groups : int
Number of groups
"""
if in_channels % groups != 0 or out_channels % groups != 0:
raise ValueError(f'Number of groups ({groups}) is not compatible with input channels '
f'({in_channels}) and/or output channels ({out_channels}), check that '
f'they are divisible by groups')
[docs]
class ConvDepth(Conv):
"""
Constructs a depthwise convolutional layer.
Supports 1D, 2D, and 3D convolution.
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,
*,
filters: int | None = None,
layer: int | None = None,
factor: float | None = None,
kernel: int | list[int] = 3,
stride: int | list[int] = 1,
padding: int | Literal['same'] | list[int] = 0,
dropout: float = 0,
activation: str | None = 'ELU',
norm: Literal[None, 'batch', 'layer'] = None,
padding_mode: Literal[None, 'zeros', 'reflect', 'replicate', 'circular'] = None,
**kwargs: Any) -> None:
"""
Parameters
----------
net_out : list[int]
Shape of the network's output, required only if layer contains factor
shapes : Shapes
Shape of the outputs from each layer
filters : int | None, Optional
Number of convolutional filters, will be used if provided, else factor will 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
Number of convolutional filters equal to the output channels multiplied by factor,
won't be used if filters is provided
kernel : int | list[int], Optional
Size of the kernel, default = 3
stride : int | list[int], Optional
Stride of the kernel, default = 1
padding : int | {'same'} | list[int], Optional
Input padding, can an int, list of ints or 'same' where 'same' preserves the input
shape, default = 0
dropout : float, Optional
Probability of dropout, default = 0
activation : str | None, Optional
Which activation function to use, default = 'ELU'
norm : {None, 'batch', 'layer'}
If batch or layer normalisation should be used
padding_mode : {None, 'zeros', 'reflect', 'replicate', 'circular'}
Padding mode to use from PyTorch, if None zero padding is used if version is >3.9.4 else
replication padding is used
**kwargs
Leftover parameters to pass to base layer for checking
"""
super().__init__(
net_out=net_out,
shapes=shapes,
filters=filters,
layer=layer,
factor=factor,
groups=shapes[-1][0],
kernel=kernel,
stride=stride,
padding=padding,
dropout=dropout,
activation=activation,
norm=norm,
padding_mode=padding_mode,
**kwargs,
)
def __getstate__(self) -> dict[str, Any]:
state: dict[str, Any] = super().__getstate__()
state.pop('groups', None)
return state
[docs]
class ConvDepthDownscale(Conv):
"""
Constructs depth downscaler using convolution with kernel size of 1.
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,
*,
dropout: float = 0,
activation: str | None = 'ELU',
norm: Literal[None, 'batch', 'layer'] = None,
padding_mode: Literal[None, 'zeros', 'reflect', 'replicate', 'circular'] = None,
**kwargs: Any) -> None:
"""
Parameters
----------
net_out : list[int]
Shape of the network's output, required only if layer contains factor
shapes: Shapes
Shape of the outputs from each layer
dropout : float, Optional
Probability of dropout, default = 0
activation : str | None, Optional
Which activation function to use, default = 'ELU'
norm : {None, 'batch', 'layer'}
If batch or layer normalisation should be used
padding_mode : {None, 'zeros', 'reflect', 'replicate', 'circular'}
Padding mode to use from PyTorch, if None zero padding is used if version is >3.9.4 else
replication padding is used
**kwargs
Leftover parameters to pass to base layer for checking
"""
super().__init__(
net_out=net_out,
shapes=shapes,
filters=1,
stride=1,
kernel=1,
padding='same',
dropout=dropout,
activation=activation,
norm=norm,
padding_mode=padding_mode,
**kwargs,
)
def __getstate__(self) -> dict[str, Any]:
state: dict[str, Any] = super().__getstate__()
state.pop('filters', None)
state.pop('stride', None)
state.pop('kernel', None)
state.pop('padding', None)
return state
[docs]
class ConvDownscale(Conv):
"""
Constructs a strided convolutional layer for downscaling.
The scale factor is equal to the stride and kernel size.
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,
*,
filters: int | None = None,
layer: int | None = None,
factor: float | None = None,
scale: int = 2,
dropout: float = 0,
activation: str | None = 'ELU',
norm: Literal[None, 'batch', 'layer'] = None,
**kwargs: Any) -> None:
"""
Parameters
----------
net_out : list[int]
Shape of the network's output, required only if layer contains factor
shapes: Shapes
Shape of the outputs from each layer
filters : int | None, Optional
Number of convolutional filters, will be used if provided, else factor will 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
Number of convolutional filters equal to the output channels multiplied by factor,
won't be used if filters is provided
scale : int, Optional
Stride and size of the kernel, which acts as the downscaling factor, default = 2
dropout : float, Optional
Probability of dropout, default = 0
activation : str | None, Optional
Which activation function to use, default = 'ELU'
norm : {None, 'batch', 'layer'}
If batch or layer normalisation should be used
**kwargs
Leftover parameters to pass to base layer for checking
"""
super().__init__(
net_out=net_out,
shapes=shapes,
filters=filters,
layer=layer,
factor=factor,
kernel=scale,
stride=scale,
padding=0,
dropout=dropout,
activation=activation,
norm=norm,
**kwargs,
)
def __getstate__(self) -> dict[str, Any]:
state: dict[str, Any] = super().__getstate__()
state.pop('padding', None)
state['scale'] = state.pop('stride')
state.pop('kernel', None)
return state
[docs]
class ConvTranspose(BaseSingleLayer):
"""
Constructs a transpose convolutional layer with fractional stride for input upscaling.
Supports 1D, 2D, and 3D transposed convolution.
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,
*,
filters: int | None = None,
layer: int | None = None,
factor: float | None = None,
kernel: int | list[int] = 3,
stride: int | list[int] = 1,
out_padding: int | list[int] = 0,
dilation: int | list[int] = 1,
padding: int | Literal['same'] | list[int] = 0,
dropout: float = 0,
padding_mode: Literal['zeros', 'reflect', 'replicate', 'circular'] = 'zeros',
activation: str | None = 'ELU',
norm: Literal[None, 'batch', 'layer'] = None,
**kwargs: Any) -> None:
"""
Parameters
----------
net_out : list[int]
Shape of the network's output, required only if layer contains factor
shapes: Shapes
Shape of the outputs from each layer
filters : int | None, Optional
Number of convolutional filters, will be used if provided, else factor will 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
Number of convolutional filters equal to the output channels multiplied by factor,
won't be used if filters is provided
kernel : int | list[int], Optional
Size of the kernel, default = 3
stride : int | list[int], Optional
Stride of the kernel, default = 1
out_padding : int | list[int], Optional
Padding applied to the output, default = 0
dilation : int | list[int], Optional
Spacing between kernel points, default = 1
padding : int | {'same'} | list[int], Optional
Inverse of convolutional padding which removes rows from each dimension in the output,
default = 0
dropout : float, Optional
Probability of dropout, default = 0
padding_mode : {'zeros', 'reflect', 'replicate', 'circular'}
Padding mode to use from PyTorch
activation : str | None, Optional
Which activation function to use, default = 'ELU'
norm : {None, 'batch', 'layer'}
If batch or layer normalisation should be used
**kwargs
Leftover parameters to pass to base layer for checking
"""
super().__init__(**kwargs)
self._slice: np.ndarray = np.array([slice(None)] * len(shapes[-1][1:]))
padding_: int | str | list[int] = padding
shape: list[int] = shapes[-1].copy()
target: list[int] = shapes[layer] if layer is not None else net_out
transpose: Type[nn.Module]
dropout_: Type[nn.Module]
batch_norm_: Type[nn.Module]
# Check for errors and calculate same padding
if isinstance(padding, str):
self._check_options('padding', padding, {'same'})
padding = _padding_transpose(kernel, stride, dilation, shapes[-1], shape)
# Check for errors and calculate number of filters
self._check_shape((2, 4), shape)
self._check_out_padding(stride, dilation, out_padding)
self._check_factor_filters(shape, filters=filters, factor=factor, target=target)
self._check_options('norm', norm, {None, 'batch', 'layer'})
transpose, dropout_, batch_norm_ = [
[nn.ConvTranspose1d, nn.Dropout1d, nn.BatchNorm1d],
[nn.ConvTranspose2d, nn.Dropout2d, nn.BatchNorm2d],
[nn.ConvTranspose3d, nn.Dropout3d, nn.BatchNorm3d],
][len(shape) - 2]
shape = _kernel_transpose_shape(
kernel,
stride,
padding,
dilation,
out_padding,
shape,
)
# Correct same padding for one-sided padding
if padding_ == 'same' and shape != shapes[-1]:
self._slice[np.array(shape[1:]) - np.array(shapes[-1][1:]) == 1] = slice(-1)
shape[1:] = shapes[-1][1:]
self.layers.add_module('Transpose', transpose(
in_channels=shapes[-1][0],
out_channels=shape[0],
kernel_size=kernel,
stride=stride,
padding=padding,
output_padding=out_padding,
padding_mode=padding_mode,
dilation=dilation,
))
# Optional layers
if activation:
self.layers.add_module('Activation', getattr(nn, activation)(
**{'inplace': True} if 'inplace' in signature(getattr(nn, activation)).parameters
else {},
))
if norm == 'batch':
self.layers.add_module('BatchNorm', batch_norm_(shape[0]))
elif norm == 'layer':
self.layers.add_module('LayerNorm', LayerNorm(shape=shape[0:1]))
if dropout:
self.layers.add_module('Dropout', dropout_(dropout))
shapes.append(shape)
def __getstate__(self) -> dict[str, Any]:
return super().__getstate__() | {
'filters': self.layers.Transpose.out_channels, # type: ignore[union-attr]
'kernel': self.layers.Transpose.kernel_size, # type: ignore[union-attr]
'stride': self.layers.Transpose.stride,
'out_padding': self.layers.Transpose.output_padding, # type: ignore[union-attr]
'dilation': self.layers.Transpose.dilation, # type: ignore[union-attr]
'padding': self.layers.Transpose.padding, # type: ignore[union-attr]
'dropout': self.layers.Dropout.p if hasattr(self.layers, 'Dropout') else 0, # type: ignore[union-attr]
'padding_mode': self.layers.Transpose.padding_mode, # type: ignore[union-attr]
'activation': self.layers.Activation.__class__.__name__
if hasattr(self.layers, 'Activation') else None,
'norm': 'batch' if hasattr(self.layers, 'BatchNorm') else
'layer' if hasattr(self.layers, 'LayerNorm') else None,
}
@staticmethod
def _check_out_padding(
stride: int | list[int],
dilation: int | list[int],
out_padding: int | list[int]) -> None:
"""
Checks if the output padding is compatible with the dilation and stride
Parameters
----------
stride : int | list[int]
Stride of the kernel
dilation : int | list[int]
Dilation of the kernel
out_padding : int | list[int]
Output padding
"""
if ((np.array(out_padding) >= np.array(stride)) *
(np.array(out_padding) >= np.array(dilation))).any():
raise ValueError(f'Output padding ({out_padding}) must be smaller than either stride '
f'({stride}) or dilation ({dilation})')
[docs]
def forward(self, x: Tensor, *_: Any, **__: Any) -> Tensor:
"""
Forward pass of the transposed convolutional layer
Parameters
----------
x : Tensor
Input tensor with shape (N,...) and type float, where N is the batch size
Returns
-------
Tensor
Output tensor with shape (N,...) and type float
"""
x = super().forward(x)
return x[..., *self._slice]
[docs]
class ConvTransposeUpscale(ConvTranspose):
"""
Constructs an upscaler using a transposed convolutional layer.
Supports 1D, 2D, and 3D transposed convolutional upscaling.
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,
*,
filters: int | None = None,
layer: int | None = None,
factor: float | None = None,
scale: int | list[int] = 2,
out_padding: int | list[int] = 0,
dropout: float = 0,
activation: str | None = 'ELU',
norm: Literal[None, 'batch', 'layer'] = None,
**kwargs: Any) -> None:
"""
Parameters
----------
net_out : list[int]
Shape of the network's output, required only if layer contains factor
shapes: Shapes
Shape of the outputs from each layer
filters : int | None, Optional
Number of convolutional filters, will be used if provided, else factor will 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
Number of convolutional filters equal to the output channels multiplied by factor,
won't be used if filters is provided
scale : int | list[int], Optional
Stride and size of the kernel, which acts as the upscaling factor, default = 2
out_padding : int | list[int], Optional
Padding applied to the output, default = 0
dropout : float, Optional
Probability of dropout, default = 0
activation : str | None, Optional
Which activation function to use, default = 'ELU'
norm : {None, 'batch', 'layer'}
If batch or layer normalisation should be used
**kwargs
Leftover parameters to pass to base layer for checking
"""
super().__init__(
net_out=net_out,
shapes=shapes,
filters=filters,
layer=layer,
factor=factor,
kernel=scale,
stride=scale,
padding=0,
out_padding=out_padding,
dropout=dropout,
activation=activation,
norm=norm,
**kwargs,
)
def __getstate__(self) -> dict[str, Any]:
state: dict[str, Any] = super().__getstate__()
state.pop('padding', None)
return state
[docs]
class ConvUpscale(Conv):
"""
Constructs an upscaler using a convolutional layer and pixel shuffling.
Supports 1D, 2D, and 3D convolutional upscaling.
See `Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel
Convolutional Neural Network <https://arxiv.org/abs/1609.05158>`_ by Shi et al. (2016) for
details.
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,
*,
filters: int | None = None,
layer: int | None = None,
factor: float | None = None,
scale: int = 2,
kernel: int | list[int] = 3,
dropout: float = 0,
activation: str | None = 'ELU',
norm: Literal[None, 'batch', 'layer'] = None,
padding_mode: Literal[None, 'zeros', 'reflect', 'replicate', 'circular'] = None,
**kwargs: Any) -> None:
"""
Parameters
----------
net_out : list[int]
Shape of the network's output, required only if layer contains factor
shapes: Shapes
Shape of the outputs from each layer
filters : int | None, Optional
Number of convolutional filters, will be used if provided, else factor will 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
Number of convolutional filters equal to the output channels multiplied by factor,
won't be used if filters is provided
scale : int, Optional
Factor to upscale the input by, default = 2
kernel : int | list[int], Optional
Size of the kernel, default = 3
dropout : float, Optional
Probability of dropout, default = 0
activation : str | None, Optional
Which activation function to use, default = 'ELU'
norm : {None, 'batch', 'layer'}
If batch or layer normalisation should be used
padding_mode : {None, 'zeros', 'reflect', 'replicate', 'circular'}
Padding mode to use from PyTorch, if None zero padding is used if version is >3.9.4 else
replicate padding is used
**kwargs
Leftover parameters to pass to base layer for checking
"""
filters_scale: int = scale ** (len(shapes[-1]) - 1)
filters = self._check_factor_filters(
shapes[-1].copy(),
filters=filters,
factor=factor,
target=shapes[layer] if layer is not None else net_out,
)[0] * filters_scale
# Convolutional layer
super().__init__(
net_out=net_out,
shapes=shapes,
filters=filters,
kernel=kernel,
stride=1,
padding='same',
dropout=dropout,
padding_mode=padding_mode,
activation=activation,
norm=norm,
**kwargs,
)
# Upscaling done using pixel shuffling
self.layers.add_module('PixelShuffle', PixelShuffle(scale))
shapes[-1][0] = shapes[-1][0] // filters_scale
shapes[-1][1:] = [length * scale for length in shapes[-1][1:]]
def __getstate__(self) -> dict[str, Any]:
state: dict[str, Any] = super().__getstate__()
state.pop('stride', None)
state.pop('padding', None)
return state
[docs]
class PixelShuffle(BaseLayer):
r"""
Used for upscaling by scale factor :math:`r` for an input :math:`(N,C\times r^n,D_1,...,D_n)` to
an output :math:`(N,C,D_1\times r,...,D_n\times r)`.
Equivalent to :class:`torch.nn.PixelShuffle`, but for nD.
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, scale: int, *, shapes: Shapes | None = None, **kwargs: Any) -> None:
"""
Parameters
----------
scale : int
Upscaling factor
shapes: Shapes | None, Optional
Shape of the outputs from each layer
**kwargs
Leftover parameters to pass to base layer for checking
"""
super().__init__(**({'idx': 0} | kwargs))
self._scale: int = scale
filters_scale: int
# If not used as an individual layer in Network
if not shapes:
return
filters_scale = self._scale ** (len(shapes[-1][1:]))
self._check_filters(filters_scale, shapes[-1])
shapes.append(shapes[-1].copy())
shapes[-1][0] = shapes[-1][0] // filters_scale
shapes[-1][1:] = [length * self._scale for length in shapes[-1][1:]]
def __getstate__(self) -> dict[str, Any]:
return super().__getstate__() | {'scale': self._scale}
@staticmethod
def _check_filters(filters_scale: int, shape: list[int]) -> None:
"""
Checks if the number of channels is an integer multiple of the upscaling factor
Parameters
----------
filters_scale : int
Upscaling factor for the number of filters
shape : list[int]
Shape of the input
"""
if shape[0] % filters_scale != 0:
raise ValueError(f'Channels ({shape}) must be an integer multiple of '
f'{filters_scale}')
[docs]
def forward(self, x: Tensor, *_: Any, **__: Any) -> Tensor:
r"""
Forward pass of pixel shuffle
Parameters
----------
x : Tensor
Input tensor with shape :math:`(N,C\times r^n,D_1,...,D_n)` and type float, where N is
the batch size, C is the number of channels, r is the upscaling factor and :math:`D_n`
is the length of dimension n
Returns
-------
Tensor
Output tensor with shape :math:`(N,C,D_1\times r,...,D_n\times r)` and type float
"""
dims: int
filters_scale: int = self._scale ** (len(x.shape[2:]))
output_channels: int = x.size(1) // filters_scale
output_shape: Tensor = self._scale * torch.tensor(x.shape[2:])
idxs: Tensor
dims = len(output_shape)
idxs = torch.arange(dims * 2) + 2
x = x.view([x.size(0), output_channels, *[self._scale] * len(x.shape[2:]), *x.shape[2:]])
x = x.permute(0, 1, *torch.ravel(torch.column_stack((idxs[dims:], idxs[:dims]))))
x = x.reshape(x.size(0), output_channels, *output_shape)
return x
def _kernel_transpose_shape(
kernel: int | list[int],
strides: int | list[int],
padding: int | list[int],
dilation: int | list[int],
out_padding: int | list[int],
shape: list[int]) -> list[int]:
"""
Calculates the output shape after a transposed convolutional kernel
Parameters
----------
kernel : int | list[int]
Size of the kernel
strides : int | list[int]
Stride of the kernel
padding : int | list[int]
Input padding
dilation : int | list[int]
Spacing between kernel elements
shape : list[int]
Input shape of the layer
Returns
-------
list[int]
Output shape of the layer
"""
shape = shape.copy()
strides, kernel, padding, dilation, out_padding = _int_list_conversion(
len(shape[1:]),
[strides, kernel, padding, dilation, out_padding]
)
for i, (stride, kernel_length, pad, dilation_length, out_pad, length) in enumerate(zip(
strides,
kernel,
padding,
dilation,
out_padding,
shape[1:]
)):
shape[i + 1] = max(
1,
stride * (length - 1) + dilation_length * (kernel_length - 1) - 2 * pad + out_pad + 1
)
return shape
def _padding_transpose(
kernel: int | list[int],
strides: int | list[int],
dilation: int | list[int],
in_shape: list[int],
out_shape: list[int]) -> list[int]:
"""
Calculates the padding required for specific output shape.
Parameters
----------
kernel : int | list[int]
Size of the kernel
strides : int | list[int]
Stride of the kernel
dilation : int | list[int]
Spacing between kernel elements
in_shape : list[int]
Input shape of the layer
out_shape : list[int]
Output shape of the layer
Returns
-------
list[int]
Required padding for specific output shape
"""
padding: list[int] = []
strides, kernel, dilation = _int_list_conversion(
len(in_shape[1:]),
[strides, kernel, dilation],
)
for stride, kernel_length, dilation_length, in_length, out_length in zip(
strides,
kernel,
dilation,
in_shape[1:],
out_shape[1:],
):
padding.append(
(stride * (in_length - 1) + dilation_length * (kernel_length - 1) - out_length + 1) // 2
)
return padding
__all__ = [
'Conv',
'ConvDepth',
'ConvDepthDownscale',
'ConvDownscale',
'ConvTranspose',
'ConvTransposeUpscale',
'ConvUpscale',
'PixelShuffle',
]