class

CvTForImageClassification

extendsPretrainedModelClassificationHeadMixin

CvTForImageClassification(config: CvTConfig)

source edit

CvT with a linear classification head (Wu et al., 2021).

Wraps the same three-stage trunk as CvT and adds a global mean pool over tokens, a final LayerNorm, and a single nn.Linear classification head:

\text{logits} = W_{\text{head}}\, \mathrm{LN}\!\bigl(\mathrm{Mean}(z^{L})\bigr) + b_{\text{head}}.

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

Parameters

configCvTConfig

Architecture specification. Must set num_classes to the desired number of output categories. See CvTConfig.

Attributes

stagesnn.ModuleList

Three CvT stages.

head_normnn.LayerNorm

Final LayerNorm applied to the mean-pooled feature.

classifiernn.Linear

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

Notes

Reference: Wu et al., "CvT: Introducing Convolutions to Vision Transformers", ICCV 2021. CvT-13 / CvT-21 reach 81.6% / 82.5% top-1 on ImageNet-1k at 224x224 (Table 1 of the paper).

Examples

End-to-end inference with the default CvT-13 classifier:
>>> import lucid
>>> from lucid.models.vision.cvt import (
...     CvTConfig, CvTForImageClassification,
... )
>>> model = CvTForImageClassification(CvTConfig())
>>> x = lucid.randn(1, 3, 224, 224)
>>> out = model(x)
>>> out.logits.shape
(1, 1000)

Used by 2

Constructors

dunder

init

→None

__init__(config: CvTConfig)

source edit

Instance methods

forward

→ImageClassificationOutput

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

source edit

End-to-end inference with the default CvT-13 classifier: >>> import lucid >>> from lucid.models.vision.cvt import ( ... CvTConfig, CvTForImageClassification, ... ) >>> model = CvTForImageClassification(CvTConfig()) >>> x = lucid.randn(1, 3, 224, 224) >>> out = model(x) >>> out.logits.shape (1, 1000)