"""
Composite layer blocks as used in networks such as ConvNeXt
"""
from typing import Any, OrderedDict, cast
from torch import nn, Tensor
from netloader.utils import Shapes
from netloader.utils.types import TensorListLike
from netloader.layers.base import BaseSingleLayer
from netloader.layers.convolutional import Conv, ConvDepth
from netloader.layers.misc import Checkpoint, DropPath, Scale
from netloader.layers import Shortcut
[docs]
class ConvNeXtBlock(BaseSingleLayer):
"""
ConvNeXt block using BaseLayer.
Original paper implementation: arXiv:2201.03545.
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 part of the parent layer to loop through in the forward pass
shapes : Shapes
Layer output shapes
check_shapes : Shapes
Checkpoint output shapes
"""
def __init__(
self,
net_out: list[int],
shapes: Shapes,
*,
drop_path: float = 0,
layer_scale: float = 1e-6,
**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
drop_path : float, Optional
Probability of Drop Path dropout, default = 0
layer_scale : float, Optional
Initial scale factor for layer Scale, default = 1e-6
**kwargs
Leftover parameters for checking if they are valid
"""
super().__init__(**kwargs)
self._drop_path: float = drop_path
self._layer_scale: float = layer_scale
self.shapes: Shapes = shapes[-1:]
self.check_shapes: Shapes = Shapes([])
self.layers = nn.Sequential(OrderedDict([
('Checkpoint', Checkpoint(self.shapes, self.check_shapes)),
('ConvDepth', ConvDepth(
net_out,
self.shapes,
layer=-1,
factor=1,
kernel=7,
stride=1,
padding=3,
activation=None,
norm='layer',
**kwargs,
)),
('Conv1', Conv(
net_out,
self.shapes,
layer=-1,
factor=4,
kernel=1,
padding='same',
activation='GELU',
**kwargs,
)),
('Conv2', Conv(
net_out,
self.shapes,
layer=0,
factor=1,
kernel=1,
padding='same',
activation=None,
**kwargs,
)),
('Scale', Scale(1, self._layer_scale, self.shapes, **kwargs)),
('DropPath', DropPath(self._drop_path, shapes=self.shapes, **kwargs)),
('Shortcut', Shortcut(
True,
-1,
self.shapes,
self.check_shapes,
**kwargs | {'checkpoint': True},
)),
]))
shapes.append(self.shapes[-1].copy())
def __getstate__(self) -> dict[str, Any]:
return super().__getstate__() | {
'drop_path': self._drop_path,
'layer_scale': self._layer_scale,
}
[docs]
def forward(
self,
x: TensorListLike,
outputs: list[TensorListLike],
checkpoints: list[TensorListLike],
*_: Any,
**__: Any) -> Tensor:
"""
Forward pass for the ConvNeXt block
Parameters
----------
x : TensorListLike
Input with tensors of shape (N,...) and type float, where N is the batch size
outputs : list[TensorListLike]
Output from each layer with tensors of shape (N,...) and type float
checkpoints : list[TensorListLike]
List of checkpoint values with tensors of shape (N,...) and type float
Returns
-------
Tensor
Output tensor with shape (N,...) and type float
"""
for layer in self.layers:
x = layer(x, outputs=outputs, checkpoints=checkpoints)
return cast(Tensor, x)
__all__ = ['ConvNeXtBlock']