class

ConvNeXtForImageClassification

extendsPretrainedModelClassificationHeadMixin

ConvNeXtForImageClassification(config: ConvNeXtConfig)

source edit

ConvNeXt with a linear classification head (Liu et al., 2022).

Wraps the same four-stage trunk as ConvNeXt (patchify stem → four stages of ConvNeXt blocks with between-stage downsamplers) and adds a global average pool, a final channel-last LayerNorm, and a single nn.Linear classification head:

\text{logits} = W_{\text{head}}\, \mathrm{LN}(\mathrm{GAP}(z^{L})) + b_{\text{head}}, \qquad W_{\text{head}} \in \mathbb{R}^{C_{\text{out}} \times d_{L}}.

Pass labels to forward to compute the cross-entropy loss in the same pass.

Parameters

configConvNeXtConfig

Architecture specification. Must set num_classes to the desired number of output categories (default 1000 for ImageNet-1k). See ConvNeXtConfig.

Attributes

stem_StemWithNorm

Patchify stem:

4 \times 4

stride-4 Conv2d + LayerNorm.

stagesnn.ModuleList

Four stages of stacked ConvNeXt blocks.

downsamplersnn.ModuleList

Three between-stage LN + stride-2 Conv2d downsamplers.

head_normnn.LayerNorm

Final LayerNorm applied to the pooled feature.

avgpoolnn.AdaptiveAvgPool2d

1 \times 1

adaptive average pool.

classifiernn.Linear

Final linear projection of width (num_classes, dims[-1]).

Notes

Reference: Zhuang Liu et al., "A ConvNet for the 2020s", CVPR 2022. ConvNeXt-T / S / B / L / XL reach 82.1 / 83.1 / 83.8 / 84.3 / 85.5 % top-1 on ImageNet-1k (Table 1 of the paper, 224x224 input) — matching or exceeding Swin at equal FLOPs.

Examples

End-to-end inference with the default ConvNeXt-T classifier:
>>> import lucid
>>> from lucid.models.vision.convnext import (
...     ConvNeXtConfig, ConvNeXtForImageClassification,
... )
>>> model = ConvNeXtForImageClassification(ConvNeXtConfig())
>>> x = lucid.randn(1, 3, 224, 224)
>>> out = model(x)
>>> out.logits.shape
(1, 1000)

Used by 2

Constructors

dunder

init

→None

__init__(config: ConvNeXtConfig)

source edit

Instance methods

forward

→ImageClassificationOutput

forward(x: Tensor, labels: Tensor | None = None)

source edit

End-to-end inference with the default ConvNeXt-T classifier: >>> import lucid >>> from lucid.models.vision.convnext import ( ... ConvNeXtConfig, ConvNeXtForImageClassification, ... ) >>> model = ConvNeXtForImageClassification(ConvNeXtConfig()) >>> x = lucid.randn(1, 3, 224, 224) >>> out = model(x) >>> out.logits.shape (1, 1000)