Vision Transformer (ViT) Implementation in PyTorch

This implementation builds a Vision Transformer from scratch. It includes patch embedding, a learnable class token and position embeddings, multi-head self-attention with fused QKV projection, an MLP feed-forward block, and a classification head.

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from einops import rearrange
from einops.layers.torch import Rearrange

class PatchEmbedding(nn.Module):
    def __init__(self, image_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_patches = (image_size // patch_size) ** 2

        # image into patches and flatten
        self.projection = nn.Sequential(
            # (batch, channel, height, width) -> (batch, num_patches, patch_dim)
            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', 
                      p1=patch_size, p2=patch_size),
            nn.Linear(patch_size * patch_size * in_channels, embed_dim)
        )

    def forward(self, x):
        return self.projection(x)
    
class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads=8, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5  

        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.attn_dropout = nn.Dropout(dropout)
        self.output_projection = nn.Linear(embed_dim, embed_dim)
        self.output_dropout = nn.Dropout(dropout)

    def forward(self, x):
        batch_size, num_tokens, embed_dim = x.shape
        
        qkv = self.qkv(x)
        qkv = qkv.reshape(batch_size, num_tokens, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, batch, heads, tokens, head_dim)
        q, k, v = qkv[0], qkv[1], qkv[2]

        # scaled dot-product attention
        attn_scores = (q @ k.transpose(-2, -1)) * self.scale
        attn_probs = attn_scores.softmax(dim=-1)
        attn_probs = self.attn_dropout(attn_probs)

        # output
        x = (attn_probs @ v).transpose(1, 2).reshape(batch_size, num_tokens, embed_dim)
        x = self.output_projection(x)
        x = self.output_dropout(x)
        return x
    
class FeedForward(nn.Module):
    def __init__(self, embed_dim, hidden_dim, dropout=0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, embed_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)
    
class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, mlp_ratio=4.0, dropout=0.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attention = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.feedforward = FeedForward(
            embed_dim,
            int(embed_dim * mlp_ratio),
            dropout
        )

    def forward(self, x):
        x = x + self.attention(self.norm1(x))
        x = x + self.feedforward(self.norm2(x))
        return x
    
class VisionTransformer(nn.Module):
    def __init__(
        self, 
        image_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.0
    ):
        super().__init__()

        self.patch_embed = PatchEmbedding(
            image_size, patch_size, in_channels, embed_dim
        )
        num_patches = self.patch_embed.num_patches

        # class token and position 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.dropout = nn.Dropout(dropout)

        self.transformer_blocks = nn.ModuleList([
            TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout)
            for _ in range(depth)
        ])

        # normalization and classification head
        self.norm = nn.LayerNorm(embed_dim)
        self.classifier = nn.Linear(embed_dim, num_classes)

        self._init_weights()
    
    def _init_weights(self):
        nn.init.normal_(self.pos_embed, std=0.02)
        nn.init.normal_(self.cls_token, std=0.02)

    def forward(self, x):
        batch_size = x.shape[0]
        
        x = self.patch_embed(x)
        
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        
        x = x + self.pos_embed
        x = self.dropout(x)
        
        for block in self.transformer_blocks:
            x = block(x)

        # take only the cls token for classification
        x = self.norm(x)
        x = x[:, 0]  
        x = self.classifier(x)
        
        return x
    
def test_vision_transformer():
    image_size = 224
    patch_size = 16
    in_channels = 3
    num_classes = 1000
    batch_size = 4
    embed_dim = 768

    # random input
    x = torch.randn(batch_size, in_channels, image_size, image_size)
    
    model = VisionTransformer(
        image_size=image_size,
        patch_size=patch_size,
        in_channels=in_channels,
        num_classes=num_classes,
        embed_dim=embed_dim,
        depth=12,
        num_heads=12,
        mlp_ratio=4.0,
        dropout=0.1
    )

    # Test forward pass
    with torch.no_grad():
        output = model(x)

    # Verify shapes
    num_patches = (image_size // patch_size) ** 2
    expected_patch_embed_shape = (batch_size, num_patches, embed_dim)
    patch_embed_output = model.patch_embed(x)
    
    print("\nVision Transformer Test Results:")
    print(f"Input shape: {x.shape}")
    print(f"Patch embedding shape: {patch_embed_output.shape}")
    print(f"Final output shape: {output.shape}")

    assert patch_embed_output.shape == expected_patch_embed_shape, \
        f"Expected patch embedding shape {expected_patch_embed_shape}, got {patch_embed_output.shape}"
    assert output.shape == (batch_size, num_classes), \
        f"Expected output shape {(batch_size, num_classes)}, got {output.shape}"
    
    print("\nAll tests passed successfully!")
    
    return model, output

if __name__ == "__main__":
    model, output = test_vision_transformer()