class

CoAtNetForImageClassification

extendsPretrainedModelClassificationHeadMixin

CoAtNetForImageClassification(config: CoAtNetConfig)

source edit

CoAtNet with a linear classification head (Dai et al., 2021).

Wraps the same conv + attention trunk as CoAtNet (stem

two MBConv stages + two relative-attention transformer stages) and adds the standard reference recipe head: global average pool → LayerNorm → optional pre-logits Linear + Tanh → linear classifier.

\text{logits} = W_{\text{cls}}\, \mathrm{Tanh}\!\bigl(W_{\text{pre}}\, \mathrm{LN}(\mathrm{GAP}(z^{S_4}))\bigr) + b_{\text{cls}}.

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

Parameters

configCoAtNetConfig

Architecture specification. Must set num_classes to the desired number of output categories. Set head_hidden_size=None to drop the pre-logits projection. See CoAtNetConfig.

Attributes

stemnn.Sequential

Two-layer stride-2 convolutional stem.

s1, s2nn.Sequential

Two MBConv stages.

s3, s4_TransformerStage

Two relative-attention transformer stages.

avgpoolnn.AdaptiveAvgPool2d

1 \times 1

adaptive average pool over spatial dims.

normnn.LayerNorm

LayerNorm applied to the pooled feature.

pre_logitsnn.Module

Either Linear + Tanh (when config.head_hidden_size is set) or an identity nn.Sequential.

classifiernn.Linear

Final linear projection of width (num_classes, head_in) where head_in is either config.head_hidden_size or config.dims[-1].

Notes

Reference: Zihang Dai et al., "CoAtNet: Marrying Convolution and Attention for All Data Sizes", NeurIPS 2021. CoAtNet-0 reaches 81.6% top-1 on ImageNet-1k at 224x224 (Table 5).

Examples

End-to-end inference with the default CoAtNet-0 classifier:
>>> import lucid
>>> from lucid.models.vision.coatnet import (
...     CoAtNetConfig, CoAtNetForImageClassification,
... )
>>> model = CoAtNetForImageClassification(CoAtNetConfig())
>>> x = lucid.randn(1, 3, 224, 224)
>>> out = model(x)
>>> out.logits.shape
(1, 1000)

Used by 2

Constructors

dunder

init

→None

__init__(config: CoAtNetConfig)

source edit

Instance methods

forward

→ImageClassificationOutput

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

source edit

End-to-end inference with the default CoAtNet-0 classifier: >>> import lucid >>> from lucid.models.vision.coatnet import ( ... CoAtNetConfig, CoAtNetForImageClassification, ... ) >>> model = CoAtNetForImageClassification(CoAtNetConfig()) >>> x = lucid.randn(1, 3, 224, 224) >>> out = model(x) >>> out.logits.shape (1, 1000)