Vision Transformer › Code Implementation
Language AI HandbookAll chapters Vision Transformer › Code ImplementationOne section · Sections in this chapter10 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement the core components of ViT to solidify our understanding. We'll build a simplified but complete version that processes images through patching and embedding before the transformer layers.
import torch
torch.manual_seed(42)We start with the patch embedding layer. This module takes an image tensor and produces a sequence of patch embeddings.
import torch.nn as nn
class PatchEmbed(nn.Module):
"""
Image to Patch Embedding.
Divides input image into non-overlapping patches and linearly projects
each flattened patch to the embedding dimension.
"""
def __init__(
self, img_size=224, patch_size=16, in_channels=3, embed_dim=768
):
super().__init__()
self.img_size = img_size
self.patch_size = patch_size
self.n_patches = (img_size // patch_size) ** 2
# Linear projection: maps P*P*C flattened patch to embed_dim
self.proj = nn.Linear(patch_size * patch_size * in_channels, embed_dim)
def forward(self, x):
# x shape: (batch_size, channels, height, width)
B, C, H, W = x.shape
# Extract non-overlapping patches using unfold
# unfold(dim, size, step) slides a window of 'size' along 'dim' with 'step'
x = x.unfold(2, self.patch_size, self.patch_size) # (B, C, H//P, P, W)
x = x.unfold(
3, self.patch_size, self.patch_size
) # (B, C, H//P, W//P, P, P)
# Permute to (B, H//P, W//P, C, P, P) then flatten spatial+channel dims
x = x.permute(0, 2, 3, 1, 4, 5)
x = x.reshape(B, self.n_patches, -1) # (B, N_patches, C*P*P)
# Project to embedding dimension: (B, N_patches, embed_dim)
x = self.proj(x)
return xNow let's test the patch embedding and inspect the output dimensions:
dummy_img = torch.randn(2, 3, 224, 224) # Batch of 2 images
patch_embed = PatchEmbed(img_size=224, patch_size=16, embed_dim=768)
patches = patch_embed(dummy_img)Input image shape: [2, 3, 224, 224] Patch embedding output: [2, 196, 768] Number of patches: 196 Patch dimensionality before projection: 768 Patch dimensionality after projection: 768
The output confirms our images are converted to sequences of 196 patch embeddings, each of dimension 768. Now let's implement the transformer building blocks. We implement multi-head attention first, returning attention weights for later visualization:
class MultiHeadAttention(nn.Module):
"""Scaled dot-product multi-head self-attention with optional attention weight output."""
def __init__(
self, dim, num_heads=8, qkv_bias=False, attn_drop=0.0, proj_drop=0.0
):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.scale = self.head_dim**-0.5
# Single linear maps all tokens to Q, K, V in one pass
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x):
B, N, C = x.shape
# Compute Q, K, V and reshape for multi-head processing
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2] # Each: (B, heads, N, head_dim)
# Scaled dot-product attention
attn = (q @ k.transpose(-2, -1)) * self.scale # (B, heads, N, N)
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
# Aggregate values and project
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x, attn # Return attention weights for visualization
class MLP(nn.Module):
"""Feed-forward network with GELU activation and 4x hidden expansion."""
def __init__(
self, in_features, hidden_features=None, out_features=None, drop=0.0
):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features * 4
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = nn.GELU()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return xThe transformer block combines these two sublayers with Pre-Norm layer normalization and residual connections:
class TransformerBlock(nn.Module):
"""Standard ViT transformer block with Pre-LayerNorm."""
def __init__(
self,
dim,
num_heads,
mlp_ratio=4.0,
qkv_bias=False,
drop=0.0,
attn_drop=0.0,
):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = MultiHeadAttention(
dim,
num_heads=num_heads,
qkv_bias=qkv_bias,
attn_drop=attn_drop,
proj_drop=drop,
)
self.norm2 = nn.LayerNorm(dim)
self.mlp = MLP(
in_features=dim, hidden_features=int(dim * mlp_ratio), drop=drop
)
def forward(self, x, return_attention=False):
# Pre-norm attention with residual: z'_l = MSA(LN(z_{l-1})) + z_{l-1}
attn_out, attn_weights = self.attn(self.norm1(x))
x = x + attn_out
# Pre-norm MLP with residual: z_l = MLP(LN(z'_l)) + z'_l
x = x + self.mlp(self.norm2(x))
if return_attention:
return x, attn_weights
return xNow the complete Vision Transformer, which assembles these components with the CLS token, position embeddings, and classification head:
class VisionTransformer(nn.Module):
"""Vision Transformer for image classification."""
def __init__(
self,
img_size=224,
patch_size=16,
in_channels=3,
num_classes=1000,
embed_dim=768,
depth=12,
num_heads=12,
mlp_ratio=4.0,
drop_rate=0.0,
):
super().__init__()
self.embed_dim = embed_dim
self.patch_embed = PatchEmbed(
img_size, patch_size, in_channels, embed_dim
)
num_patches = self.patch_embed.n_patches
# Learnable CLS token: starts as noise, trained to aggregate image info
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
# Learned 1D position embeddings: one per patch + one for the CLS token
self.pos_embed = nn.Parameter(
torch.zeros(1, num_patches + 1, embed_dim)
)
self.pos_drop = nn.Dropout(p=drop_rate)
# L identical transformer blocks
self.blocks = nn.ModuleList(
[
TransformerBlock(
dim=embed_dim,
num_heads=num_heads,
mlp_ratio=mlp_ratio,
drop=drop_rate,
)
for _ in range(depth)
]
)
# Final layer norm and classification head
self.norm = nn.LayerNorm(embed_dim)
self.head = nn.Linear(embed_dim, num_classes)
# Initialize with small normal noise to break symmetry
nn.init.normal_(self.cls_token, std=0.02)
nn.init.normal_(self.pos_embed, std=0.02)
def forward(self, x, return_attention=False):
B = x.shape[0]
# Step 1: Patch embedding
x = self.patch_embed(x) # (B, N, D)
# Step 2: Prepend CLS token
cls_tokens = self.cls_token.expand(B, -1, -1) # (B, 1, D)
x = torch.cat((cls_tokens, x), dim=1) # (B, N+1, D)
# Step 3: Add position embeddings
x = x + self.pos_embed
x = self.pos_drop(x)
# Step 4: Process through transformer layers
attn_weights = []
for block in self.blocks:
if return_attention:
x, attn = block(x, return_attention=True)
attn_weights.append(attn)
else:
x = block(x)
# Step 5: Extract CLS token and classify
x = self.norm(x)
cls_output = x[:, 0] # CLS token only: (B, D)
logits = self.head(cls_output) # (B, num_classes)
if return_attention:
return logits, attn_weights
return logitsLet's instantiate a small model suitable for demonstration, verify the output dimensions, and check the parameter count:
# Small ViT configuration for demonstration
model = VisionTransformer(
img_size=224,
patch_size=16,
embed_dim=192, # Tiny model: ViT-Tiny width
depth=4, # Shallow for fast demo
num_heads=3,
num_classes=10,
)
total_params = sum(p.numel() for p in model.parameters())
test_input = torch.randn(1, 3, 224, 224)
output = model(test_input)
sample_logits = output.squeeze()[:5]Model parameters: 1,965,130 Input shape: [1, 3, 224, 224] Output shape: [1, 10] Sample logits (first 5): ['-0.094', '0.791', '0.309', '-0.657', '-0.169']
The compact model configuration yields a manageable parameter count while maintaining all the architectural properties of full-scale Vision Transformers. Now let's visualize the core patching operation that turns a 2D image into a 1D token sequence:
The conversion from the structured 2D image on the left to the flat sequence on the right captures the core challenge: the sequence representation loses the two-dimensional adjacency that is natural in the image. Patches that were side-by-side are now adjacent in the sequence, but patches that were vertically adjacent (14 positions apart in row-major order) appear far apart. The model must learn from position embeddings and training data to reconstruct these spatial relationships.
Finally, let's visualize how self-attention distributes weights across patches:
The attention heatmap illustrates two distinct learned behaviors. The top row (CLS token) exhibits diffuse, roughly uniform attention across all patches. This reflects its role as a global aggregator that must summarize the entire image for classification. The remaining rows show concentration along the diagonal. This reflects locality: patches attend most strongly to spatially adjacent patches. No convolutional architecture enforced this locality; the model discovered it from data because local correlations are a measurable statistical property of natural images. ViT learns the same visual priors as CNNs, but through gradient descent rather than architectural design.