"""
PyTorch-Network-Loader Python implementation of ConvNeXt.
See the ConvNeXt paper: arXiv:2201.03545.
See the ConvNeXt GitHub: https://github.com/facebookresearch/ConvNeXt.
"""
from typing import Any, cast
import numpy as np
from torch import nn
from netloader import layers
from netloader.network import Network
from netloader.utils import Config, Shapes
from netloader.layers.base import BaseLayer
[docs]
class ConvNeXt(Network):
"""
ConvNeXt network using Network.
Original paper implementation: arXiv:2201.03545.
Attributes
----------
layer_num : int, Optional
Number of layers to use, if None use all layers
group : int
Which group is the active group if a layer has the group attribute
name : str
Name of the network, used for saving
version : str
NetLoader version string
checkpoints : list[TensorListLike]
Outputs from each checkpoint with each Tensor having shape (N,...) and type float, where N
is the batch size
config : dict[str, Any]
Network configuration
kl_loss : Tensor
KL divergence loss on the latent space, of shape (1) and type float, if using a sample layer
net : TypedModuleList
Network construction
shapes : Shapes
Layer output shapes
check_shapes : Shapes
Checkpoint output shapes
"""
def __init__(
self,
name: str,
in_shape: list[int] | tuple[int, ...],
out_shape: list[int],
*,
max_drop_path: float = 0,
layer_scale: float = 1e-6,
dims: list[int] | None = None,
depths: list[int] | None = None) -> None:
"""
Parameters
----------
name : str
Name of the network configuration file
in_shape : list[int] | tuple[int, ...]
Shape of the input tensor, excluding batch size
out_shape : list[int]
shape of the output tensor, excluding batch size
max_drop_path : float, Optional
Maximum drop path fraction, default = 0
layer_scale : float, Optional
Default value for learnable parameters to scale the convolutional filters in a ConvNeXt
block, default = 1e-6
dims : list[int] | None, Optional
Number of convolutional filters in each downscaled section of the network,
default = [96, 192, 384, 768]
depths : list[int] | None, Optional
Number of ConvNeXt blocks in each downscaled section of the network,
default = [3, 3, 9, 3]
"""
super().__init__(
name,
Config(),
in_shape,
[],
suppress_warning=True,
)
self._layer_scale: float = layer_scale
self._max_drop_path: float = max_drop_path
self._dims: list[int] = dims or [96, 192, 384, 768]
self._depths: list[int] = depths or [3, 3, 9, 3]
self._build_net(out_shape)
self.apply(self._init_weights)
def __getstate__(self) -> dict[str, Any]:
"""
Returns a dictionary containing the state of the network for pickling
Returns
-------
dict[str, Any]
Dictionary containing the state of the network
"""
return super().__getstate__() | {
'dims': self._dims,
'depths': self._depths,
'layer_scale': self._layer_scale,
'max_drop_path': self._max_drop_path,
}
def __setstate__(self, state: dict[str, Any]) -> None:
"""
Sets the state of the network for pickling
Parameters
----------
state : dict[str, Any]
Dictionary containing the state of the network
"""
super().__init__(
state['name'],
Config(),
state['shapes'][0],
[],
suppress_warning=True,
)
self._layer_scale = state['layer_scale']
self._max_drop_path = state['max_drop_path']
self._dims = state['dims']
self._depths = state['depths']
self.shapes = Shapes([state['shapes'][0]])
self._build_net(state['shapes'][-1])
self._load_state_dict(cast(dict[str, Any], state.get('layers', state.get('net'))))
@staticmethod
def _init_weights(layer: nn.Module) -> None:
"""
Initialises the weights of the convolutional and linear layers.
Parameters
----------
layer : Module
Layer to initialise if it is convolutional or linear
"""
if isinstance(layer, (nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.Linear)):
nn.init.trunc_normal_(layer.weight, std=0.02)
if layer.bias is not None:
nn.init.constant_(layer.bias, 0)
def _build_net(self, out_shape: list[int]) -> None:
"""
Constructs the ConvNeXt network layers.
Parameters
----------
out_shape : list[int]
shape of the output tensor, excluding batch size
"""
assert isinstance(self.shapes, list)
drop_paths: list[float]
depths: list[int] = np.cumsum(self._depths).tolist()
kwargs: dict[str, Any] = {'check_shapes': Shapes()}
drop_paths = np.linspace(0, self._max_drop_path, depths[-1]).tolist()
# Stem
self.layers.extend(self._stem(out_shape, **kwargs))
# Main body
for dim, in_depth, out_depth in zip(self._dims, [0, *depths], depths):
# Downscaling, this is skipped in the first iteration
self.layers.extend([
layers.LayerNorm(dims=1, shapes=self.shapes, **kwargs),
layers.ConvDownscale(
out_shape,
self.shapes,
filters=dim,
scale=2,
activation=None if self.version != '<3.8.1' else 'ELU',
**kwargs,
),
] if dim != self._dims[0] else [])
# ConvNeXt blocks
self.layers.extend([layers.ConvNeXtBlock(
out_shape,
self.shapes,
drop_path=drop_path,
layer_scale=self._layer_scale,
) for drop_path in drop_paths[in_depth:out_depth]])
# Head
self.layers.extend(self._head(out_shape, **kwargs))
def _head(self, out_shape: list[int], **kwargs: Any) -> list[BaseLayer]:
"""
Head of ConvNeXt for adapting the learned features into the desired output.
Parameters
----------
out_shape : list[int]
shape of the output tensor, excluding batch size
**kwargs
Global network parameters
Returns
-------
list[BaseLayer]
Layers used in the head
"""
return [
layers.AdaptivePool(1, self.shapes, **kwargs),
layers.Reshape([-1], shapes=self.shapes, **kwargs),
layers.LayerNorm(dims=1, shapes=self.shapes, **kwargs),
layers.Linear(
out_shape,
self.shapes,
factor=1,
activation=None if self.version != '<3.8.1' else 'SELU',
**kwargs,
),
]
def _stem(self, out_shape: list[int], **kwargs: Any) -> list[BaseLayer]:
"""
Stem of ConvNeXt for initial downscaling of the input.
Parameters
----------
out_shape : list[int]
shape of the output tensor, excluding batch size
**kwargs
Global network parameters
Returns
-------
list[BaseLayer]
Layers used in the stem
"""
return [
layers.ConvDownscale(
out_shape,
self.shapes,
filters=self._dims[0],
scale=4,
activation=None if self.version != '<3.8.1' else 'ELU',
norm='layer',
**kwargs,
),
]
[docs]
class ConvNeXtTiny(ConvNeXt):
"""
Tiny version of ConvNeXt with the number of convolutional filters in each downscaled section
following [96, 192, 384, 768] with the number of ConvNeXt blocks in each downscaled section
following [3, 3, 9, 3].
Original paper implementation: arXiv:2201.03545.
Attributes
----------
layer_num : int, Optional
Number of layers to use, if None use all layers
group : int
Which group is the active group if a layer has the group attribute
name : str
Name of the network, used for saving
version : str
NetLoader version string
checkpoints : list[TensorListLike]
Outputs from each checkpoint with each Tensor having shape (N,...) and type float, where N
is the batch size
config : dict[str, Any]
Network configuration
kl_loss : Tensor
KL divergence loss on the latent space, of shape (1) and type float, if using a sample layer
net : TypedModuleList
Network construction
shapes : Shapes
Layer output shapes
check_shapes : Shapes
Checkpoint output shapes
"""
_name: str = 'convnext_tiny'
def __init__(
self,
in_shape: list[int] | tuple[int, ...],
out_shape: list[int],
*,
max_drop_path: float = 0,
layer_scale: float = 1e-6) -> None:
"""
Parameters
----------
in_shape : list[int] | list[list[int]] | tuple[int, ...]
Shape of the input tensor(s), excluding batch size
out_shape : list[int]
shape of the output tensor, excluding batch size
max_drop_path : float, Optional
Maximum drop path fraction, default = 0
layer_scale : float, Optional
Default value for learnable parameters to scale the convolutional filters in a ConvNeXt
block, default = 1e-6
"""
super().__init__(
self._name,
in_shape,
out_shape,
max_drop_path=max_drop_path,
layer_scale=layer_scale,
)
[docs]
class ConvNeXtSmall(ConvNeXt):
"""
Small version of ConvNeXt with the number of convolutional filters in each downscaled section
following [96, 192, 384, 768] with the number of ConvNeXt blocks in each downscaled section
following [3, 3, 27, 3].
Original paper implementation: arXiv:2201.03545.
Attributes
----------
layer_num : int, Optional
Number of layers to use, if None use all layers
group : int
Which group is the active group if a layer has the group attribute
name : str
Name of the network, used for saving
version : str
NetLoader version string
checkpoints : list[TensorListLike]
Outputs from each checkpoint with each Tensor having shape (N,...) and type float, where N
is the batch size
config : dict[str, Any]
Network configuration
kl_loss : Tensor
KL divergence loss on the latent space, of shape (1) and type float, if using a sample layer
net : TypedModuleList
Network construction
shapes : Shapes
Layer output shapes
check_shapes : Shapes
Checkpoint output shapes
"""
_name: str = 'convnext_small'
def __init__(
self,
in_shape: list[int] | tuple[int, ...],
out_shape: list[int],
*,
max_drop_path: float = 0,
layer_scale: float = 1e-6) -> None:
"""
Parameters
----------
in_shape : list[int] | list[list[int]] | tuple[int, ...]
Shape of the input tensor(s), excluding batch size
out_shape : list[int]
shape of the output tensor, excluding batch size
max_drop_path : float, Optional
Maximum drop path fraction, default = 0
layer_scale : float, Optional
Default value for learnable parameters to scale the convolutional filters in a ConvNeXt
block, default = 1e-6
"""
super().__init__(
self._name,
in_shape,
out_shape,
max_drop_path=max_drop_path,
layer_scale=layer_scale,
depths=[3, 3, 27, 3],
)
[docs]
class ConvNeXtBase(ConvNeXt):
"""
Base version of ConvNeXt with the number of convolutional filters in each downscaled section
following [128, 256, 512, 1024] with the number of ConvNeXt blocks in each downscaled section
following [3, 3, 27, 3].
Original paper implementation: arXiv:2201.03545.
Attributes
----------
layer_num : int, Optional
Number of layers to use, if None use all layers
group : int
Which group is the active group if a layer has the group attribute
name : str
Name of the network, used for saving
version : str
NetLoader version string
checkpoints : list[TensorListLike]
Outputs from each checkpoint with each Tensor having shape (N,...) and type float, where N
is the batch size
config : dict[str, Any]
Network configuration
kl_loss : Tensor
KL divergence loss on the latent space, of shape (1) and type float, if using a sample layer
net : TypedModuleList
Network construction
shapes : Shapes
Layer output shapes
check_shapes : Shapes
Checkpoint output shapes
"""
_name: str = 'convnext_base'
def __init__(
self,
in_shape: list[int] | tuple[int, ...],
out_shape: list[int],
*,
max_drop_path: float = 0,
layer_scale: float = 1e-6) -> None:
"""
Parameters
----------
in_shape : list[int] | list[list[int]] | tuple[int, ...]
Shape of the input tensor(s), excluding batch size
out_shape : list[int]
shape of the output tensor, excluding batch size
max_drop_path : float, Optional
Maximum drop path fraction, default = 0
layer_scale : float, Optional
Default value for learnable parameters to scale the convolutional filters in a ConvNeXt
block, default = 1e-6
"""
super().__init__(
self._name,
in_shape,
out_shape,
max_drop_path=max_drop_path,
layer_scale=layer_scale,
dims=[128, 256, 512, 1024],
depths=[3, 3, 27, 3],
)
[docs]
class ConvNeXtLarge(ConvNeXt):
"""
Large version of ConvNeXt with the number of convolutional filters in each downscaled section
following [192, 384, 768, 1536] with the number of ConvNeXt blocks in each downscaled section
following [3, 3, 27, 3].
Original paper implementation: arXiv:2201.03545.
Attributes
----------
layer_num : int, Optional
Number of layers to use, if None use all layers
group : int
Which group is the active group if a layer has the group attribute
name : str
Name of the network, used for saving
version : str
NetLoader version string
checkpoints : list[TensorListLike]
Outputs from each checkpoint with each Tensor having shape (N,...) and type float, where N
is the batch size
config : dict[str, Any]
Network configuration
kl_loss : Tensor
KL divergence loss on the latent space, of shape (1) and type float, if using a sample layer
net : TypedModuleList
Network construction
shapes : Shapes
Layer output shapes
check_shapes : Shapes
Checkpoint output shapes
"""
_name: str = 'convnext_large'
def __init__(
self,
in_shape: list[int] | tuple[int, ...],
out_shape: list[int],
*,
max_drop_path: float = 0,
layer_scale: float = 1e-6) -> None:
"""
Parameters
----------
in_shape : list[int] | list[list[int]] | tuple[int, ...]
Shape of the input tensor(s), excluding batch size
out_shape : list[int]
shape of the output tensor, excluding batch size
max_drop_path : float, Optional
Maximum drop path fraction, default = 0
layer_scale : float, Optional
Default value for learnable parameters to scale the convolutional filters in a ConvNeXt
block, default = 1e-6
"""
super().__init__(
self._name,
in_shape,
out_shape,
max_drop_path=max_drop_path,
layer_scale=layer_scale,
dims=[192, 384, 768, 1536],
depths=[3, 3, 27, 3],
)
[docs]
class ConvNeXtXLarge(ConvNeXt):
"""
Extra large version of ConvNeXt with the number of convolutional filters in each downscaled
section following [256, 512, 1024, 2048] with the number of ConvNeXt blocks in each downscaled
section following [3, 3, 27, 3].
Original paper implementation: arXiv:2201.03545.
Attributes
----------
layer_num : int, Optional
Number of layers to use, if None use all layers
group : int
Which group is the active group if a layer has the group attribute
name : str
Name of the network, used for saving
version : str
NetLoader version string
checkpoints : list[TensorListLike]
Outputs from each checkpoint with each Tensor having shape (N,...) and type float, where N
is the batch size
config : dict[str, Any]
Network configuration
kl_loss : Tensor
KL divergence loss on the latent space, of shape (1) and type float, if using a sample layer
net : TypedModuleList
Network construction
shapes : Shapes
Layer output shapes
check_shapes : Shapes
Checkpoint output shapes
"""
_name: str = 'convnext_xlarge'
def __init__(
self,
in_shape: list[int] | tuple[int, ...],
out_shape: list[int],
*,
max_drop_path: float = 0,
layer_scale: float = 1e-6) -> None:
"""
Parameters
----------
in_shape : list[int] | list[list[int]] | tuple[int, ...]
Shape of the input tensor(s), excluding batch size
out_shape : list[int]
shape of the output tensor, excluding batch size
max_drop_path : float, Optional
Maximum drop path fraction, default = 0
layer_scale : float, Optional
Default value for learnable parameters to scale the convolutional filters in a ConvNeXt
block, default = 1e-6
"""
super().__init__(
self._name,
in_shape,
out_shape,
max_drop_path=max_drop_path,
layer_scale=layer_scale,
dims=[256, 512, 1024, 2048],
depths=[3, 3, 27, 3],
)
__all__ = [
'ConvNeXt',
'ConvNeXtTiny',
'ConvNeXtSmall',
'ConvNeXtBase',
'ConvNeXtLarge',
'ConvNeXtXLarge',
]