data

ViTConfig

extendsModelConfig

ViTConfig(image_size: int = 224, patch_size: int = 16, num_classes: int = 1000, in_channels: int = 3, dim: int = 768, depth: int = 12, num_heads: int = 12, mlp_ratio: float = 4.0, dropout: float = 0.0, attention_dropout: float = 0.0, layer_norm_eps: float = 1e-06)

source edit

Configuration dataclass for every Vision Transformer (ViT) variant.

ViTConfig is an immutable container that fully specifies the architecture of a Vision Transformer. It is consumed by both ViT (backbone) and ViTForImageClassification (classifier head on top of the CLS token). All canonical variants described in Dosovitskiy et al. (2020) — Base / Large / Huge at patch sizes 14, 16, and 32 — can be expressed by choosing different values for dim, depth, num_heads, and patch_size.

The number of input tokens fed into the transformer encoder is $N + 1$ , where $N = (\text{image\_size} / \text{patch\_size})^2$ is the number of non-overlapping patches and the extra token is a learnable [CLS] embedding prepended to the sequence.

Parameters

image_sizeint= 224

Spatial side length (in pixels) of the square input image. image_size must be divisible by patch_size. Defaults to 224, matching the ImageNet recipe in the original paper.

patch_sizeint= 16

Side length of the non-overlapping square patches extracted by the patch-embedding convolution. Common values are 16, 32 (Base / Large) or 14 (Huge). Defaults to 16.

num_classesint= 1000

Number of output logits produced by the classification head. Only used by ViTForImageClassification. Defaults to 1000 (ImageNet-1k).

in_channelsint= 3

Number of input image channels. Defaults to 3 (RGB).

dimint= 768

Hidden width

D

of every transformer block — also the dimensionality of the patch embedding and of the CLS token. Defaults to 768 (ViT-Base).

depthint= 12

Number of stacked transformer encoder blocks

L

. Defaults to 12 (ViT-Base).

num_headsint= 12

Number of attention heads inside each multi-head self-attention block. dim must be divisible by num_heads. Defaults to 12.

mlp_ratiofloat= 4.0

Expansion ratio of the MLP block. The hidden width of the MLP is int(dim * mlp_ratio); the standard recipe is 4.0 (i.e. a 4x expansion). Defaults to 4.0.

dropoutfloat= 0.0

Dropout probability applied (i) after the patch embedding + positional embedding sum and (ii) inside the MLP feed-forward sub-layer. Defaults to 0.0.

attention_dropoutfloat= 0.0

Dropout probability applied to the post-softmax attention weights inside every multi-head self-attention block. Defaults to 0.0.

layer_norm_epsfloat= 1e-06

Epsilon added to the variance inside every nn.LayerNorm (the two per-block norms and the final pre-head norm). The original ViT formulation fixes this at 1e-6. Defaults to 1e-6.

Attributes

model_typeClassVar[str]

Constant string "vit" used by the model registry to associate config instances with the ViT model class.

Notes

The canonical variants registered as factory functions in lucid.models.vision.vit are (see Table 1 of the paper):

============== ===== ====== ======= =========== Variant dim depth heads patch sizes ============== ===== ====== ======= =========== ViT-Base 768 12 12 16, 32 ViT-Large 1024 24 16 16, 32 ViT-Huge 1280 32 16 14 ============== ===== ====== ======= ===========

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 custom ViT-Base/16 configuration for CIFAR-10 (10 classes,
32x32 inputs, patch size 4 — yields 64 tokens):
>>> from lucid.models.vision.vit import ViTConfig
>>> cfg = ViTConfig(
...     image_size=32,
...     patch_size=4,
...     num_classes=10,
...     dim=768,
...     depth=12,
...     num_heads=12,
... )
>>> cfg.image_size, cfg.patch_size, cfg.num_classes
(32, 4, 10)
Reuse a canonical variant and override only the number of classes:
>>> from lucid.models.vision.vit._pretrained import _CFG_B16
>>> from dataclasses import replace
>>> cfg = replace(_CFG_B16, num_classes=100)
>>> cfg.dim, cfg.depth, cfg.num_classes
(768, 12, 100)

Used by 3

Constructors

dunder

init

→None

__init__(image_size: int = 224, patch_size: int = 16, num_classes: int = 1000, in_channels: int = 3, dim: int = 768, depth: int = 12, num_heads: int = 12, mlp_ratio: float = 4.0, dropout: float = 0.0, attention_dropout: float = 0.0, layer_norm_eps: float = 1e-06)

source edit

ViTConfig(image_size: int = 224, patch_size: int = 16, num_classes: int = 1000, in_channels: int = 3, dim: int = 768, depth: int = 12, num_heads: int = 12, mlp_ratio: float = 4.0, dropout: float = 0.0, attention_dropout: float = 0.0, layer_norm_eps: float = 1e-06)

Build a custom ViT-Base/16 configuration for CIFAR-10 (10 classes, 32x32 inputs, patch size 4 — yields 64 tokens): >>> from lucid.models.vision.vit import ViTConfig >>> cfg = ViTConfig( ... image_size=32, ... patch_size=4, ... num_classes=10, ... dim=768, ... depth=12, ... num_heads=12, ... ) >>> cfg.image_size, cfg.patch_size, cfg.num_classes (32, 4, 10) Reuse a canonical variant and override only the number of classes: >>> from lucid.models.vision.vit._pretrained import _CFG_B16 >>> from dataclasses import replace >>> cfg = replace(_CFG_B16, num_classes=100) >>> cfg.dim, cfg.depth, cfg.num_classes (768, 12, 100)

__init__(image_size: int = 224, patch_size: int = 16, num_classes: int = 1000, in_channels: int = 3, dim: int = 768, depth: int = 12, num_heads: int = 12, mlp_ratio: float = 4.0, dropout: float = 0.0, attention_dropout: float = 0.0, layer_norm_eps: float = 1e-06)