class

ViT

extendsPretrainedModelBackboneMixin

ViT(config: ViTConfig)

source edit

Vision Transformer backbone (Dosovitskiy et al., 2020).

A ViT first splits the input image into a sequence of non-overlapping patches and linearly projects each patch into a $D$ -dimensional embedding. A learnable [CLS] token is prepended, a learnable positional embedding is added, and the resulting sequence is processed by depth stacked pre-norm transformer encoder blocks. The final representation returned by forward_features is the CLS token embedding after a final LayerNorm — the standard feature used for classification or transfer learning.

Concretely, the input sequence fed to the encoder is

z_0 = [\,x_{\text{class}};\, x_p^1 E;\, x_p^2 E;\, \dots;\, x_p^N E\,] + E_{\text{pos}},

where $x_p^i \in \mathbb{R}^{P^2 C}$ is the $i$ -th flattened patch, $E \in \mathbb{R}^{P^2 C \times D}$ is the patch projection, $E_{\text{pos}} \in \mathbb{R}^{(N+1) \times D}$ is the learnable positional embedding, and $N = (\text{image\_size} / \text{patch\_size})^2$ .

Use this backbone when you want features rather than logits — e.g. fine-tuning on a custom downstream head, contrastive pretraining, or feature extraction for retrieval. For end-to-end image classification, use ViTForImageClassification instead.

Parameters

configViTConfig

Frozen dataclass specifying image_size, patch_size, dim, depth, num_heads, mlp_ratio, dropout, attention_dropout, and in_channels. See ViTConfig.

Attributes

patch_embed_PatchEmbed

Patch-extraction convolution producing the initial token sequence.

cls_tokenTensor

Learnable [CLS] token, shape (1, 1, dim). Initialized with truncated normal

\mathcal{N}(0, 0.02^2)

pos_embedTensor

Learnable positional embedding, shape (1, num_patches + 1, dim). Initialized with truncated normal

\mathcal{N}(0, 0.02^2)

pos_dropnn.Dropout

Dropout applied right after adding positional embeddings.

blocksnn.ModuleList

config.depth stacked _ViTBlock encoder layers.

normnn.LayerNorm

Final LayerNorm applied before extracting the CLS token.

feature_infolist[FeatureInfo]

Single-stage feature description with channel count equal to config.dim and spatial reduction equal to config.patch_size.

Notes

The classification head is not included in this class — only the backbone trunk. forward_features returns a flat $(B, D)$ feature tensor, whereas forward wraps it in a BaseModelOutput whose last_hidden_state has shape $(B, 1, D)$ for spatial consistency with other vision models.

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

Examples

Build a ViT-Base/16 backbone and run a single forward pass:
>>> import lucid
>>> from lucid.models.vision.vit import ViT, ViTConfig
>>> cfg = ViTConfig()              # ViT-Base/16 defaults
>>> model = ViT(cfg)
>>> x = lucid.randn(1, 3, 224, 224)
>>> feat = model.forward_features(x)
>>> feat.shape                     # (B, dim)
(1, 768)
>>> out = model(x)
>>> out.last_hidden_state.shape    # (B, 1, dim)
(1, 1, 768)

Used by 2

Constructors

dunder

init

→None

__init__(config: ViTConfig)

source edit

Initialise the Flatten module. See the class docstring for parameter semantics.

Properties

prop

feature_info

→list[FeatureInfo]

feature_info: list[FeatureInfo]

source edit

Instance methods

forward

→BaseModelOutput

forward(x: Tensor)

source edit

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

forward_features

→Tensor

forward_features(x: Tensor)

source edit

Build a ViT-Base/16 backbone and run a single forward pass: >>> import lucid >>> from lucid.models.vision.vit import ViT, ViTConfig >>> cfg = ViTConfig() # ViT-Base/16 defaults >>> model = ViT(cfg) >>> x = lucid.randn(1, 3, 224, 224) >>> feat = model.forward_features(x) >>> feat.shape # (B, dim) (1, 768) >>> out = model(x) >>> out.last_hidden_state.shape # (B, 1, dim) (1, 1, 768)