Vision Transformers: The Revolution in Image Processing

Vision Transformers (ViT) introduced a new approach to image processing, using the Transformer architecture originally designed for natural language processing.

Vision Transformer Architecture

1. Patch Embedding

class PatchEmbed(nn.Module):
    def __init__(self, img_size, patch_size, 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
        
        self.proj = nn.Conv2d(
            in_channels,
            embed_dim,
            kernel_size=patch_size,
            stride=patch_size
        )
        
    def forward(self, x):
        x = self.proj(x)  # (B, E, H', W')
        x = x.flatten(2)  # (B, E, N)
        x = x.transpose(1, 2)  # (B, N, E)
        return x

2. Vision Transformer

class VisionTransformer(nn.Module):
    def __init__(self, img_size, patch_size, in_channels, n_classes, embed_dim, 
                 num_heads, num_layers, mlp_ratio=4., qkv_bias=True):
        super().__init__()
        
        self.patch_embed = PatchEmbed(img_size, patch_size, in_channels, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_embed.n_patches + 1, embed_dim))
        self.pos_drop = nn.Dropout(p=0.1)
        
        self.blocks = nn.ModuleList([
            TransformerBlock(embed_dim, num_heads, mlp_ratio, qkv_bias)
            for _ in range(num_layers)
        ])
        
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, n_classes)
        
        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.shape[0]
        
        # Patch embedding
        x = self.patch_embed(x)
        
        # Add the classification token
        cls_token = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_token, x), dim=1)
        
        # Position embedding
        x = x + self.pos_embed
        x = self.pos_drop(x)
        
        # Transformer blocks
        for block in self.blocks:
            x = block(x)
            
        x = self.norm(x)
        x = x[:, 0]  # Use the classification token
        
        return self.head(x)

Practical Applications

1. Image Classification

from transformers import ViTImageProcessor, ViTForImageClassification

# Load the model and the processor
processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')

# Prediction
image = Image.open('image.jpg')
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)

2. Object Detection

class DETR(nn.Module):
    def __init__(self, num_classes, num_queries, hidden_dim):
        super().__init__()
        
        self.backbone = VisionTransformer(...)
        self.transformer = Transformer(...)
        self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        
    def forward(self, x):
        features = self.backbone(x)
        hs = self.transformer(features)
        
        outputs_class = self.class_embed(hs)
        outputs_coord = self.bbox_embed(hs).sigmoid()
        
        return {'pred_logits': outputs_class[-1], 
                'pred_boxes': outputs_coord[-1]}

Advantages Over CNNs

1. Global Attention

   • Better understanding of the overall context
   • Captures long-range relationships
   • More flexible than local convolutions

2. Flexibility

   • Uniform architecture for different tasks
   • Easily extensible
   • Better adaptation to the data

3. Performance

   • Superior results on certain tasks
   • Better generalization
   • More efficient on large datasets

Limitations and Solutions

1. Data Requirements

• Solution: Pre-training on large datasets
• Fine-tuning on specific tasks

2. Computational Complexity

• Solution: Efficient patch embedding
• Attention optimization

3. Memory

• Solution: Sparse attention
• Gradient checkpointing

Usage Examples

Medical Image Classification

# Configuration for medical images
model = VisionTransformer(
    img_size=224,
    patch_size=16,
    in_channels=3,
    n_classes=num_classes,
    embed_dim=768,
    num_heads=12,
    num_layers=12
)

# Training
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

for epoch in range(num_epochs):
    model.train()
    for batch in train_loader:
        images, labels = batch
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    scheduler.step()

Conclusion

Vision Transformers represent a major advance in image processing, offering a powerful alternative to traditional CNNs. Their flexible architecture and ability to capture global relationships make them a valuable tool for many applications.

Additional Resources

• ViT Paper
• DETR Paper
• Hugging Face Vision Transformers