Lately, I’ve been fascinated by how transformers have transformed computer vision. What started as a breakthrough in natural language processing now reshapes how we approach image recognition. I remember when convolutional neural networks dominated every computer vision task, but Vision Transformers (ViTs) have challenged that status quo. Their ability to capture long-range dependencies across entire images intrigued me enough to build one from scratch in PyTorch. If you’ve ever wondered how these architectures actually work under the hood, join me in this practical exploration where we’ll implement every component ourselves.

The core idea behind ViTs is surprisingly straightforward. Instead of processing images through convolutional filters, we split them into fixed-size patches. Think of these patches as visual words. Each patch gets flattened into a vector, creating a sequence of tokens similar to how text transformers process sentences. We then add positional information since spatial relationships matter in images—unlike text where word order carries meaning. This sequence becomes input for a standard transformer encoder, which applies multi-head self-attention to understand relationships between patches. Finally, a classification head predicts the image category. Why does this approach outperform traditional CNNs in certain scenarios? It comes down to the model’s ability to weigh all parts of the image simultaneously.

Let’s start by setting up our environment. First, ensure you have Python 3.8+ installed. Create a virtual environment and install these essentials:

pip install torch torchvision torchaudio
pip install matplotlib seaborn timm wandb

Now, import critical libraries:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10
import matplotlib.pyplot as plt

The heart of our ViT is the patch embedding layer. This converts image patches into trainable vectors. Notice how we use a convolutional layer for efficiency—even though we’re building a transformer, clever borrowing from CNNs helps:

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.projection = nn.Conv2d(in_channels, embed_dim, 
                                   kernel_size=patch_size, 
                                   stride=patch_size)
        
    def forward(self, x):
        x = self.projection(x)  # [batch, embed_dim, grid, grid]
        x = x.flatten(2).transpose(1, 2)  # [batch, num_patches, embed_dim]
        return x

Multi-head self-attention is where the magic happens. How do we make the model focus on relevant patches? By computing attention scores between every patch pair:

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12):
        super().__init__()
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5
        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.proj = nn.Linear(embed_dim, embed_dim)
        
    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(2)  # Separate query, key, value
        
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        return self.proj(x)

For training, we’ll use CIFAR-10. Preprocessing is crucial—ViTs need correctly sized inputs:

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225])
])
train_data = CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)

During training, I discovered some techniques that significantly boost accuracy. Layer scaling stabilizes learning, especially with deeper architectures. Mixup augmentation—blending images and labels—improves generalization. Here’s how I implement it:

def mixup_data(x, y, alpha=0.2):
    lam = np.random.beta(alpha, alpha) if alpha > 0 else 1
    index = torch.randperm(x.size(0))
    mixed_x = lam * x + (1 - lam) * x[index]
    return mixed_x, y, y[index], lam

After training, visualizing attention maps reveals what the model learns. Notice how it focuses on distinctive features—like a cat’s ears or a car’s wheels:

def visualize_attention(model, img):
    model.eval()
    with torch.no_grad():
        attn = model.get_last_attention(img.unsqueeze(0))
    plt.imshow(attn[0, 0].cpu().numpy())  # First head

For deployment, consider converting your model to TorchScript. This creates a serialized version that runs without Python dependencies:

scripted_model = torch.jit.script(model)
scripted_model.save("vit_scripted.pt")

Building ViTs from scratch taught me more than any pre-trained model ever could. You understand every design choice and its trade-offs. Try experimenting with different patch sizes—how would using 32x32 patches instead of 16x16 affect performance? Or adjust the number of transformer blocks? If you found this guide helpful, share it with others diving into modern computer vision. Have questions or insights? Let’s discuss in the comments—I’d love to hear about your implementation experiences!