Vision Transformers in Computer Vision

The Story That Explains Vision Transformers

A Vision Transformer (ViT) applies the Transformer architecture — originally invented for text — directly to images. Instead of words, it splits an image into small fixed-size patches, treats each patch as a "token", and lets every patch attend to every other patch simultaneously through self-attention. The result is a model that understands global context from the very first layer, something CNNs only achieve deep in their stacks.

A Brief History — From Pixels to Patches

To appreciate ViTs, you need to understand what came before — and why it wasn't enough.

CNN vs ViT — The Fundamental Difference

How ViT Works — Step by Step

Let us trace a single image through a Vision Transformer from pixel array to class prediction.

Self-Attention — The Heart of ViT

Formally, self-attention computes three vectors for each patch token:

Multi-Head Attention

Instead of one set of Q, K, V projections, ViT uses multiple heads in parallel — typically 12 or 16. Each head learns to attend to a different aspect: one head may focus on local texture, another on long-range spatial relationships, another on colour continuity. Outputs from all heads are concatenated and projected back to dimension D.

ViT Architecture Variants — The Family Tree

The original paper introduced three model sizes. The community has since expanded this family dramatically.

Swin Transformer — The Windowed Attention Trick

Positional Encoding — Teaching ViT Where Things Are

A Transformer is permutation-invariant by design. Shuffle all the patches randomly and without positional encoding, the model cannot tell the difference. For images — where a cat's ears are above its nose, and the sky is above the ground — spatial position is crucial.

Code — Building a ViT from Scratch in PyTorch

Let us implement a minimal but complete Vision Transformer. Every line maps directly to the theory above.

import torch
import torch.nn as nn
import math

# ─────────────────────────────────────────────────────
#  1. PATCH EMBEDDING
#     Splits image into patches and linearly projects each
# ─────────────────────────────────────────────────────
class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.num_patches = (img_size // patch_size) ** 2   # 196
        # Conv2d with kernel=patch_size, stride=patch_size = non-overlapping patches
        self.proj = nn.Conv2d(in_channels, embed_dim,
                               kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        x = self.proj(x)         # (B, embed_dim, H/P, W/P)
        x = x.flatten(2)         # (B, embed_dim, num_patches)
        x = x.transpose(1, 2)   # (B, num_patches, embed_dim)
        return x

# ─────────────────────────────────────────────────────
#  2. MULTI-HEAD SELF-ATTENTION
# ─────────────────────────────────────────────────────
class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim  = embed_dim // num_heads          # 64
        self.scale     = self.head_dim ** -0.5           # 1/√64
        self.qkv  = nn.Linear(embed_dim, embed_dim * 3)  # project to Q, K, V at once
        self.proj = nn.Linear(embed_dim, embed_dim)
        self.attn_drop = nn.Dropout(dropout)

    def forward(self, x):
        B, N, C = x.shape
        # Split into Q, K, V and reshape for multi-head
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, heads, N, head_dim)
        q, k, v = qkv.unbind(0)

        # 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)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        return self.proj(x)

# ─────────────────────────────────────────────────────
#  3. TRANSFORMER ENCODER BLOCK
# ─────────────────────────────────────────────────────
class TransformerBlock(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4.0, dropout=0.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn  = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
        self.norm2 = nn.LayerNorm(embed_dim)
        mlp_hidden = int(embed_dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, mlp_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_hidden, embed_dim),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        x = x + self.attn(self.norm1(x))   # Pre-norm + residual
        x = x + self.mlp(self.norm2(x))    # Pre-norm + residual
        return x

# ─────────────────────────────────────────────────────
#  4. FULL VISION TRANSFORMER
# ─────────────────────────────────────────────────────
class VisionTransformer(nn.Module):
    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, dropout=0.1):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
        num_patches = self.patch_embed.num_patches

        # Learnable [CLS] token and positional embeddings
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.pos_drop  = nn.Dropout(dropout)

        # Stack of Transformer encoder blocks
        self.blocks = nn.Sequential(*[
            TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout)
            for _ in range(depth)
        ])

        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)

        # Initialise weights
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)

    def forward(self, x):
        B = x.size(0)
        x = self.patch_embed(x)                          # (B, 196, 768)

        cls = self.cls_token.expand(B, -1, -1)            # (B, 1, 768)
        x = torch.cat((cls, x), dim=1)                 # (B, 197, 768)
        x = self.pos_drop(x + self.pos_embed)

        x = self.blocks(x)                              # 12 Transformer layers
        x = self.norm(x)

        cls_out = x[:, 0]                               # Extract [CLS] token
        return self.head(cls_out)                        # (B, num_classes)

# ─── Instantiate and verify ────────────────────────
model = VisionTransformer(
    img_size=224, patch_size=16, embed_dim=768, depth=12,
    num_heads=12, num_classes=1000
)
dummy = torch.randn(2, 3, 224, 224)  # batch of 2 RGB images
out   = model(dummy)
print(f"Output shape: {out.shape}")
params = sum(p.numel() for p in model.parameters())
print(f"Parameters:   {params:,}")

Code — Fine-Tuning a Pretrained ViT with Hugging Face

Training ViT from scratch requires massive datasets. In practice, you fine-tune a pretrained model. Hugging Face transformers makes this three lines of model code.

from transformers import ViTForImageClassification, ViTImageProcessor
from datasets import load_dataset
import torch
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

# ── 1. Load pretrained ViT-Base/16 (ImageNet-21k checkpoint) ──
MODEL_NAME = "google/vit-base-patch16-224-in21k"
num_classes = 10  # e.g. CIFAR-10

processor = ViTImageProcessor.from_pretrained(MODEL_NAME)
model = ViTForImageClassification.from_pretrained(
    MODEL_NAME,
    num_labels=num_classes,
    ignore_mismatched_sizes=True   # replaces the 21k classifier head
)

# ── 2. Load and preprocess dataset ──────────────────────────
dataset = load_dataset("cifar10")

def transform_fn(examples):
    inputs = processor(images=examples["img"], return_tensors="pt")
    inputs["labels"] = torch.tensor(examples["label"])
    return inputs

dataset = dataset.with_transform(transform_fn)
train_loader = DataLoader(dataset["train"], batch_size=32, shuffle=True)
val_loader   = DataLoader(dataset["test"],  batch_size=64)

# ── 3. Training loop ─────────────────────────────────────────
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

optimizer = AdamW(model.parameters(), lr=2e-4, weight_decay=0.01)
scheduler = CosineAnnealingLR(optimizer, T_max=10)

for epoch in range(10):
    model.train()
    total_loss = 0
    for batch in train_loader:
        pixel_values = batch["pixel_values"].to(device)
        labels        = batch["labels"].to(device)
        outputs = model(pixel_values=pixel_values, labels=labels)
        loss    = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        total_loss += loss.item()

    # Validation accuracy
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for batch in val_loader:
            pv  = batch["pixel_values"].to(device)
            lbl = batch["labels"].to(device)
            out = model(pixel_values=pv)
            pred = out.logits.argmax(-1)
            correct += (pred == lbl).sum().item()
            total   += lbl.size(0)

    scheduler.step()
    print(f"Epoch {epoch+1:2d} | Loss: {total_loss/len(train_loader):.4f} | "
          f"Val Acc: {correct/total*100:.2f}%")

Visualising Attention Maps

One of ViT's unique advantages: attention weights are directly interpretable. We can visualise what patches each token attended to, giving genuine insight into what the model "sees".

import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from transformers import ViTModel, ViTImageProcessor

# Load model with output_attentions=True
processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
model     = ViTModel.from_pretrained("google/vit-base-patch16-224",
                                      output_attentions=True)
model.eval()

# Preprocess a test image
img    = Image.open("test_image.jpg").convert("RGB")
inputs = processor(images=img, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

# attentions: list of (batch, heads, N+1, N+1) per layer
# Take last layer, average across heads, focus on [CLS] → patch attention
last_layer_attn = outputs.attentions[-1]          # (1, 12, 197, 197)
avg_head_attn   = last_layer_attn[0].mean(dim=0)  # (197, 197)
cls_attn        = avg_head_attn[0, 1:]            # [CLS] attending to 196 patches
attn_map        = cls_attn.reshape(14, 14).numpy()

# Upscale attention map to image resolution
attn_resized = np.array(Image.fromarray(attn_map).resize(
    (224, 224), Image.BILINEAR))

# Plot original + attention overlay
fig, axes = plt.subplots(1, 2, figsize=(10, 5))
axes[0].imshow(img); axes[0].set_title("Original Image")
axes[1].imshow(img)
axes[1].imshow(attn_resized, alpha=0.6, cmap="jet")
axes[1].set_title("Attention Map — [CLS] token")
for ax in axes: ax.axis("off")
plt.tight_layout()
plt.savefig("attention_map.png", dpi=150)
plt.show()

Masked Autoencoders (MAE) — Self-Supervised ViT Pretraining

ViT for Dense Prediction — Object Detection & Segmentation

Original ViT produces a single [CLS] vector — perfect for classification, useless for detection or segmentation, which require per-pixel or per-region features. Several approaches solve this.

ViT vs CNN — When to Use Which

Training Tips & Tricks for ViTs

ViTs require different training recipes than CNNs. These are the non-negotiable adjustments.

Complete Training Recipe — Code

import torch
import torch.nn as nn
from torchvision import transforms, datasets
from timm import create_model
from timm.data.mixup import Mixup
from timm.loss import SoftTargetCrossEntropy
from timm.scheduler import CosineLRScheduler
from torch.optim import AdamW

# ── 1. Model — ViT-Base/16 from timm ──────────────────────
model = create_model(
    "vit_base_patch16_224",
    pretrained=True,
    num_classes=10
)

# ── 2. Augmentation pipeline (DeiT-style) ─────────────────
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.RandAugment(num_ops=2, magnitude=9),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                           [0.229, 0.224, 0.225]),
    transforms.RandomErasing(p=0.25),
])

# ── 3. Mixup + CutMix (timm's Mixup handles both) ─────────
mixup_fn = Mixup(
    mixup_alpha=0.8, cutmix_alpha=1.0,
    prob=1.0, switch_prob=0.5,
    mode="batch", label_smoothing=0.1, num_classes=10
)
criterion = SoftTargetCrossEntropy()   # works with soft Mixup labels

# ── 4. Optimiser — split weight decay groups ───────────────
no_decay = ["bias", "LayerNorm.weight"]
optim_groups = [
    {"params": [p for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)],
     "weight_decay": 0.05},
    {"params": [p for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)],
     "weight_decay": 0.0},
]
optimizer = AdamW(optim_groups, lr=1e-3)

# ── 5. Cosine LR schedule with warm-up ────────────────────
scheduler = CosineLRScheduler(
    optimizer, t_initial=100,      # 100 total epochs
    lr_min=1e-5,
    warmup_t=5,                    # 5-epoch warm-up
    warmup_lr_init=1e-6,
)

# ── 6. Training loop ───────────────────────────────────────
device = torch.device("cuda")
model.to(device)

for epoch in range(100):
    model.train()
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        images, soft_labels = mixup_fn(images, labels)
        logits = model(images)
        loss   = criterion(logits, soft_labels)
        loss.backward()
        # Gradient clipping — important for ViT stability
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad()
    scheduler.step(epoch + 1)

ViT Performance on Major Benchmarks

Compute & Memory Complexity — The Numbers

Real-World Applications

The Future — What Comes After ViT?

Quick Reference — Key ViT Parameters