Skip to content

ConvNeXt

The ConvNeXt model is based on the "A ConvNet for the 2020s" paper.

Architecture overview

This architecture compiles tricks from transformer-based vision models to improve a pure convolutional model.

ConvNeXt architecture

The key takeaways from the paper are the following:

  • update the stem convolution to act like a patchify layer of transformers
  • increase block kernel size to 7
  • switch to depth-wise convolutions
  • reduce the amount of activations and normalization layers

Model builders

The following model builders can be used to instantiate a ConvNeXt model, with or without pre-trained weights. All the model builders internally rely on the ConvNeXt base class.

ConvNeXt

ConvNeXt(num_blocks: list[int], planes: list[int], num_classes: int = 10, in_channels: int = 3, conv_layer: Callable[..., Module] | None = None, act_layer: Module | None = None, norm_layer: Callable[[int], Module] | None = None, drop_layer: Callable[..., Module] | None = None, stochastic_depth_prob: float = 0.0)

Bases: Sequential

Source code in holocron/models/classification/convnext.py
def __init__(
    self,
    num_blocks: list[int],
    planes: list[int],
    num_classes: int = 10,
    in_channels: int = 3,
    conv_layer: Callable[..., nn.Module] | None = None,
    act_layer: nn.Module | None = None,
    norm_layer: Callable[[int], nn.Module] | None = None,
    drop_layer: Callable[..., nn.Module] | None = None,
    stochastic_depth_prob: float = 0.0,
) -> None:
    if conv_layer is None:
        conv_layer = nn.Conv2d
    if norm_layer is None:
        norm_layer = partial(LayerNorm2d, eps=1e-6)
    if act_layer is None:
        act_layer = nn.GELU()
    self.dilation: int = 1

    # Patchify-like stem
    layers = conv_sequence(
        in_channels,
        planes[0],
        None,
        norm_layer,
        drop_layer,
        conv_layer,
        kernel_size=4,
        stride=4,
        padding=0,
        bias=True,
    )

    block_idx = 0
    tot_blocks = sum(num_blocks)
    for _num_blocks, _planes, _oplanes in zip(num_blocks, planes, [*planes[1:], planes[-1]], strict=True):
        # adjust stochastic depth probability based on the depth of the stage block
        sd_probs = [stochastic_depth_prob * (block_idx + _idx) / (tot_blocks - 1.0) for _idx in range(_num_blocks)]
        stage: list[nn.Module] = [
            Bottlenext(_planes, act_layer, norm_layer, drop_layer, stochastic_depth_prob=sd_prob)
            for _idx, sd_prob in zip(range(_num_blocks), sd_probs, strict=True)
        ]
        if _planes != _oplanes:
            stage.append(
                nn.Sequential(
                    LayerNorm2d(_planes),
                    nn.Conv2d(_planes, _oplanes, kernel_size=2, stride=2),
                )
            )
        layers.append(nn.Sequential(*stage))
        block_idx += _num_blocks

    super().__init__(
        OrderedDict([
            ("features", nn.Sequential(*layers)),
            ("pool", GlobalAvgPool2d(flatten=True)),
            (
                "head",
                nn.Sequential(
                    nn.LayerNorm(planes[-1], eps=1e-6),
                    nn.Linear(planes[-1], num_classes),
                ),
            ),
        ])
    )

    # Init all layers
    for m in self.modules():
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            nn.init.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)

convnext_atto

convnext_atto(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> ConvNeXt

ConvNeXt-Atto variant of Ross Wightman inspired by "A ConvNet for the 2020s"

PARAMETER DESCRIPTION
pretrained

If True, returns a model pre-trained on ImageNette

TYPE: bool DEFAULT: False

checkpoint

If specified, the model's parameters will be set to the checkpoint's values

TYPE: Checkpoint | None DEFAULT: None

progress

If True, displays a progress bar of the download to stderr

TYPE: bool DEFAULT: True

kwargs

keyword args of ConvNeXt

TYPE: Any DEFAULT: {}

RETURNS DESCRIPTION
ConvNeXt

classification model

ConvNeXt_Atto_Checkpoint

Bases: Enum

IMAGENETTE class-attribute instance-attribute
IMAGENETTE = _checkpoint(arch='convnext_atto', url='https://github.com/frgfm/Holocron/releases/download/v0.2.1/convnext_atto_224-f38217e7.pth', acc1=0.8759, acc5=0.9832, sha256='f38217e7361060e6fe00e8fa95b0e8774150190eed9e55c812bbd3b6ab378ce9', size=13535258, num_params=3377730, commit='d4a59999179b42fc0d3058ac6b76cc41f49dd56e', train_args='./imagenette2-320/ --arch convnext_atto --batch-size 64 --mixup-alpha 0.2 --amp --device 0 --epochs 100 --lr 1e-3 --label-smoothing 0.1 --random-erase 0.1 --train-crop-size 176 --val-resize-size 232 --opt adamw --weight-decay 5e-2')
DEFAULT class-attribute instance-attribute
DEFAULT = IMAGENETTE
Source code in holocron/models/classification/convnext.py
def convnext_atto(
    pretrained: bool = False,
    checkpoint: Checkpoint | None = None,
    progress: bool = True,
    **kwargs: Any,
) -> ConvNeXt:
    """ConvNeXt-Atto variant of Ross Wightman inspired by
    ["A ConvNet for the 2020s"](https://arxiv.org/pdf/2201.03545.pdf)

    Args:
        pretrained: If True, returns a model pre-trained on ImageNette
        checkpoint: If specified, the model's parameters will be set to the checkpoint's values
        progress: If True, displays a progress bar of the download to stderr
        kwargs: keyword args of [`ConvNeXt`][holocron.models.classification.convnext.ConvNeXt]

    Returns:
        classification model

    ::: holocron.models.ConvNeXt_Atto_Checkpoint
        options:
            heading_level: 4
            show_if_no_docstring: true
    """
    checkpoint = _handle_legacy_pretrained(
        pretrained,
        checkpoint,
        ConvNeXt_Atto_Checkpoint.DEFAULT.value,
    )
    return _convnext(checkpoint, progress, [2, 2, 6, 2], [40, 80, 160, 320], **kwargs)

convnext_femto

convnext_femto(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> ConvNeXt

ConvNeXt-Femto variant of Ross Wightman inspired by "A ConvNet for the 2020s"

PARAMETER DESCRIPTION
pretrained

If True, returns a model pre-trained on ImageNette

TYPE: bool DEFAULT: False

checkpoint

If specified, the model's parameters will be set to the checkpoint's values

TYPE: Checkpoint | None DEFAULT: None

progress

If True, displays a progress bar of the download to stderr

TYPE: bool DEFAULT: True

kwargs

keyword args of ConvNeXt

TYPE: Any DEFAULT: {}

RETURNS DESCRIPTION
ConvNeXt

classification model

Source code in holocron/models/classification/convnext.py
def convnext_femto(
    pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any
) -> ConvNeXt:
    """ConvNeXt-Femto variant of Ross Wightman inspired by
    ["A ConvNet for the 2020s"](https://arxiv.org/pdf/2201.03545.pdf)

    Args:
        pretrained: If True, returns a model pre-trained on ImageNette
        checkpoint: If specified, the model's parameters will be set to the checkpoint's values
        progress: If True, displays a progress bar of the download to stderr
        kwargs: keyword args of [`ConvNeXt`][holocron.models.classification.convnext.ConvNeXt]

    Returns:
        classification model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _convnext(checkpoint, progress, [2, 2, 6, 2], [48, 96, 192, 384], **kwargs)

convnext_pico

convnext_pico(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> ConvNeXt

ConvNeXt-Pico variant of Ross Wightman inspired by "A ConvNet for the 2020s"

PARAMETER DESCRIPTION
pretrained

If True, returns a model pre-trained on ImageNette

TYPE: bool DEFAULT: False

checkpoint

If specified, the model's parameters will be set to the checkpoint's values

TYPE: Checkpoint | None DEFAULT: None

progress

If True, displays a progress bar of the download to stderr

TYPE: bool DEFAULT: True

kwargs

keyword args of ConvNeXt

TYPE: Any DEFAULT: {}

RETURNS DESCRIPTION
ConvNeXt

classification model

Source code in holocron/models/classification/convnext.py
def convnext_pico(
    pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any
) -> ConvNeXt:
    """ConvNeXt-Pico variant of Ross Wightman inspired by
    ["A ConvNet for the 2020s"](https://arxiv.org/pdf/2201.03545.pdf)

    Args:
        pretrained: If True, returns a model pre-trained on ImageNette
        checkpoint: If specified, the model's parameters will be set to the checkpoint's values
        progress: If True, displays a progress bar of the download to stderr
        kwargs: keyword args of [`ConvNeXt`][holocron.models.classification.convnext.ConvNeXt]

    Returns:
        classification model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _convnext(checkpoint, progress, [2, 2, 6, 2], [64, 128, 256, 512], **kwargs)

convnext_nano

convnext_nano(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> ConvNeXt

ConvNeXt-Nano variant of Ross Wightman inspired by "A ConvNet for the 2020s"

PARAMETER DESCRIPTION
pretrained

If True, returns a model pre-trained on ImageNette

TYPE: bool DEFAULT: False

checkpoint

If specified, the model's parameters will be set to the checkpoint's values

TYPE: Checkpoint | None DEFAULT: None

progress

If True, displays a progress bar of the download to stderr

TYPE: bool DEFAULT: True

kwargs

keyword args of ConvNeXt

TYPE: Any DEFAULT: {}

RETURNS DESCRIPTION
ConvNeXt

classification model

Source code in holocron/models/classification/convnext.py
def convnext_nano(
    pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any
) -> ConvNeXt:
    """ConvNeXt-Nano variant of Ross Wightman inspired by
    ["A ConvNet for the 2020s"](https://arxiv.org/pdf/2201.03545.pdf)

    Args:
        pretrained: If True, returns a model pre-trained on ImageNette
        checkpoint: If specified, the model's parameters will be set to the checkpoint's values
        progress: If True, displays a progress bar of the download to stderr
        kwargs: keyword args of [`ConvNeXt`][holocron.models.classification.convnext.ConvNeXt]

    Returns:
        classification model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _convnext(checkpoint, progress, [2, 2, 8, 2], [80, 160, 320, 640], **kwargs)

convnext_tiny

convnext_tiny(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> ConvNeXt

ConvNeXt-T from "A ConvNet for the 2020s"

PARAMETER DESCRIPTION
pretrained

If True, returns a model pre-trained on ImageNette

TYPE: bool DEFAULT: False

checkpoint

If specified, the model's parameters will be set to the checkpoint's values

TYPE: Checkpoint | None DEFAULT: None

progress

If True, displays a progress bar of the download to stderr

TYPE: bool DEFAULT: True

kwargs

keyword args of ConvNeXt

TYPE: Any DEFAULT: {}

RETURNS DESCRIPTION
ConvNeXt

classification model

Source code in holocron/models/classification/convnext.py
def convnext_tiny(
    pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any
) -> ConvNeXt:
    """ConvNeXt-T from
    ["A ConvNet for the 2020s"](https://arxiv.org/pdf/2201.03545.pdf)

    Args:
        pretrained: If True, returns a model pre-trained on ImageNette
        checkpoint: If specified, the model's parameters will be set to the checkpoint's values
        progress: If True, displays a progress bar of the download to stderr
        kwargs: keyword args of [`ConvNeXt`][holocron.models.classification.convnext.ConvNeXt]

    Returns:
        classification model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _convnext(checkpoint, progress, [3, 3, 9, 3], [96, 192, 384, 768], **kwargs)

convnext_small

convnext_small(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> ConvNeXt

ConvNeXt-S from "A ConvNet for the 2020s"

PARAMETER DESCRIPTION
pretrained

If True, returns a model pre-trained on ImageNette

TYPE: bool DEFAULT: False

checkpoint

If specified, the model's parameters will be set to the checkpoint's values

TYPE: Checkpoint | None DEFAULT: None

progress

If True, displays a progress bar of the download to stderr

TYPE: bool DEFAULT: True

kwargs

keyword args of ConvNeXt

TYPE: Any DEFAULT: {}

RETURNS DESCRIPTION
ConvNeXt

classification model

Source code in holocron/models/classification/convnext.py
def convnext_small(
    pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any
) -> ConvNeXt:
    """ConvNeXt-S from
    ["A ConvNet for the 2020s"](https://arxiv.org/pdf/2201.03545.pdf)

    Args:
        pretrained: If True, returns a model pre-trained on ImageNette
        checkpoint: If specified, the model's parameters will be set to the checkpoint's values
        progress: If True, displays a progress bar of the download to stderr
        kwargs: keyword args of [`ConvNeXt`][holocron.models.classification.convnext.ConvNeXt]

    Returns:
        classification model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _convnext(checkpoint, progress, [3, 3, 27, 3], [96, 192, 384, 768], **kwargs)

convnext_base

convnext_base(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> ConvNeXt

ConvNeXt-B from "A ConvNet for the 2020s"

PARAMETER DESCRIPTION
pretrained

If True, returns a model pre-trained on ImageNette

TYPE: bool DEFAULT: False

checkpoint

If specified, the model's parameters will be set to the checkpoint's values

TYPE: Checkpoint | None DEFAULT: None

progress

If True, displays a progress bar of the download to stderr

TYPE: bool DEFAULT: True

kwargs

keyword args of ConvNeXt

TYPE: Any DEFAULT: {}

RETURNS DESCRIPTION
ConvNeXt

classification model

Source code in holocron/models/classification/convnext.py
def convnext_base(
    pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any
) -> ConvNeXt:
    """ConvNeXt-B from
    ["A ConvNet for the 2020s"](https://arxiv.org/pdf/2201.03545.pdf)

    Args:
        pretrained: If True, returns a model pre-trained on ImageNette
        checkpoint: If specified, the model's parameters will be set to the checkpoint's values
        progress: If True, displays a progress bar of the download to stderr
        kwargs: keyword args of [`ConvNeXt`][holocron.models.classification.convnext.ConvNeXt]

    Returns:
        classification model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _convnext(checkpoint, progress, [3, 3, 27, 3], [128, 256, 512, 1024], **kwargs)

convnext_large

convnext_large(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> ConvNeXt

ConvNeXt-L from "A ConvNet for the 2020s"

PARAMETER DESCRIPTION
pretrained

If True, returns a model pre-trained on ImageNette

TYPE: bool DEFAULT: False

checkpoint

If specified, the model's parameters will be set to the checkpoint's values

TYPE: Checkpoint | None DEFAULT: None

progress

If True, displays a progress bar of the download to stderr

TYPE: bool DEFAULT: True

kwargs

keyword args of ConvNeXt

TYPE: Any DEFAULT: {}

RETURNS DESCRIPTION
ConvNeXt

classification model

Source code in holocron/models/classification/convnext.py
def convnext_large(
    pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any
) -> ConvNeXt:
    """ConvNeXt-L from
    ["A ConvNet for the 2020s"](https://arxiv.org/pdf/2201.03545.pdf)

    Args:
        pretrained: If True, returns a model pre-trained on ImageNette
        checkpoint: If specified, the model's parameters will be set to the checkpoint's values
        progress: If True, displays a progress bar of the download to stderr
        kwargs: keyword args of [`ConvNeXt`][holocron.models.classification.convnext.ConvNeXt]

    Returns:
        classification model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _convnext(checkpoint, progress, [3, 3, 27, 3], [192, 384, 768, 1536], **kwargs)

convnext_xl

convnext_xl(pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any) -> ConvNeXt

ConvNeXt-XL from "A ConvNet for the 2020s"

PARAMETER DESCRIPTION
pretrained

If True, returns a model pre-trained on ImageNette

TYPE: bool DEFAULT: False

checkpoint

If specified, the model's parameters will be set to the checkpoint's values

TYPE: Checkpoint | None DEFAULT: None

progress

If True, displays a progress bar of the download to stderr

TYPE: bool DEFAULT: True

kwargs

keyword args of ConvNeXt

TYPE: Any DEFAULT: {}

RETURNS DESCRIPTION
ConvNeXt

classification model

Source code in holocron/models/classification/convnext.py
def convnext_xl(
    pretrained: bool = False, checkpoint: Checkpoint | None = None, progress: bool = True, **kwargs: Any
) -> ConvNeXt:
    """ConvNeXt-XL from
    ["A ConvNet for the 2020s"](https://arxiv.org/pdf/2201.03545.pdf)

    Args:
        pretrained: If True, returns a model pre-trained on ImageNette
        checkpoint: If specified, the model's parameters will be set to the checkpoint's values
        progress: If True, displays a progress bar of the download to stderr
        kwargs: keyword args of [`ConvNeXt`][holocron.models.classification.convnext.ConvNeXt]

    Returns:
        classification model
    """
    checkpoint = _handle_legacy_pretrained(pretrained, checkpoint, None)
    return _convnext(checkpoint, progress, [3, 3, 27, 3], [256, 512, 1024, 2048], **kwargs)