class

ViTForImageClassification

extendsPretrainedModelClassificationHeadMixin

ViTForImageClassification(config: ViTConfig)

source edit

Vision Transformer with a linear classification head on the CLS token.

Wraps the same trunk as ViT (patch embedding -> positional embedding -> depth transformer blocks -> final LayerNorm) and adds a single nn.Linear classification head that maps the final CLS-token feature of dimension $D$ to config.num_classes logits:

\text{logits} = W_{\text{head}}\, z_L^{[0]} + b_{\text{head}}, \qquad W_{\text{head}} \in \mathbb{R}^{C \times D}

where $z_L^{[0]}$ is the post-LayerNorm CLS token at the final layer and $C$ is the number of classes. This is the standard ImageNet recipe described in the original ViT paper.

Pass labels to forward to compute the cross-entropy loss in the same pass — handy for HuggingFace-style training loops.

Parameters

configViTConfig

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

Attributes

patch_embed_PatchEmbed

Patch-extraction convolution producing the initial token sequence.

cls_tokenTensor

Learnable CLS token, shape (1, 1, dim).

pos_embedTensor

Learnable positional embedding, shape (1, num_patches + 1, dim).

pos_dropnn.Dropout

Dropout applied right after adding positional embeddings.

blocksnn.ModuleList

Stacked transformer encoder blocks.

normnn.LayerNorm

Final LayerNorm applied before extracting the CLS token.

headnn.Linear

Classification head with shape (num_classes, dim). The attribute is intentionally named head so pretrained state-dicts using the standard head.weight / head.bias key names load directly.

Notes

Reference: Alexey Dosovitskiy et al., "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale", ICLR 2021, arXiv:2010.11929.

Examples

End-to-end inference with the default ViT-Base/16 classifier:
>>> import lucid
>>> from lucid.models.vision.vit import (
...     ViTConfig, ViTForImageClassification
... )
>>> cfg = ViTConfig(num_classes=1000)
>>> model = ViTForImageClassification(cfg)
>>> x = lucid.randn(1, 3, 224, 224)
>>> out = model(x)
>>> out.logits.shape
(1, 1000)
Pass labels for joint loss computation during training:
>>> labels = lucid.tensor([0])
>>> out = model(x, labels=labels)
>>> out.loss.shape
()

Used by 2

Constructors

dunder

init

→None

__init__(config: ViTConfig)

source edit

Instance methods

forward

→ImageClassificationOutput

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

source edit

Return bias tensor of shape (1, num_heads, ws², ws²).

End-to-end inference with the default ViT-Base/16 classifier: >>> import lucid >>> from lucid.models.vision.vit import ( ... ViTConfig, ViTForImageClassification ... ) >>> cfg = ViTConfig(num_classes=1000) >>> model = ViTForImageClassification(cfg) >>> x = lucid.randn(1, 3, 224, 224) >>> out = model(x) >>> out.logits.shape (1, 1000) Pass labels for joint loss computation during training: >>> labels = lucid.tensor([0]) >>> out = model(x, labels=labels) >>> out.loss.shape ()