Vision Transformer (ViT) - Khi Transformer Học Cách Nhìn

import torch
import torch.nn as nn
import numpy as np

class PatchExtractor(nn.Module):
    def __init__(self, patch_size=16):
        super().__init__()
        self.patch_size = patch_size

    def forward(self, x):
        # x shape: (batch_size, channels, height, width)
        B, C, H, W = x.size()
        
        assert H % self.patch_size == 0 and W % self.patch_size == 0, \
            f"Image size ({H}x{W}) phải chia hết cho patch_size ({self.patch_size})"

        num_patches_h = H // self.patch_size
        num_patches_w = W // self.patch_size
        num_patches = num_patches_h * num_patches_w

        # unfold: trích xuất sliding windows theo chiều H và W
        patches = (
            x.unfold(2, self.patch_size, self.patch_size)
             .unfold(3, self.patch_size, self.patch_size)
             .permute(0, 2, 3, 1, 4, 5)
             .contiguous()
             .view(B, num_patches, -1)
        )
        # Output shape: (batch_size, num_patches, patch_size*patch_size*channels)
        return patches

class EmbeddingLayer(nn.Module):
    def __init__(self, latent_size=1024, num_patches=4, input_dim=768):
        super().__init__()
        self.num_patches = num_patches
        
        # Chiếu patch thô → latent space
        self.input_embedder = nn.Linear(input_dim, latent_size)
        
        # Học positional encoding từ chỉ số vị trí
        self.pos_embedder = nn.Linear(1, latent_size)
        
        # Tạo chỉ số vị trí: [0, 1, 2, ..., num_patches-1]
        self.register_buffer(
            'positional_information',
            torch.arange(0, num_patches).reshape(1, num_patches, 1).float()
        )

    def forward(self, x):
        # x shape: (N, num_patches, input_dim)
        N = x.shape[0]
        
        input_embedding = self.input_embedder(x)
        
        pos_info = self.positional_information.expand(N, -1, -1)
        positional_embedding = self.pos_embedder(pos_info)
        
        # Cộng patch embedding và positional embedding
        return input_embedding + positional_embedding

class ViT(nn.Module):
    def __init__(
        self,
        patch_size: int = 16,
        img_dimension: tuple = (32, 32),
        latent_size: int = 1024,
        num_heads: int = 2,
        num_classes: int = 10,
        dropout: float = 0.1,
    ):
        super().__init__()
        
        H, W = img_dimension
        assert H % patch_size == 0 and W % patch_size == 0, \
            "Kích thước ảnh phải chia hết cho patch_size!"

        self.num_patches = (H // patch_size) * (W // patch_size)
        input_dim = patch_size * patch_size * 3  # 3 channels RGB

        # --- Các thành phần chính ---
        self.patchifier = PatchExtractor(patch_size)
        
        self.embedding_layer = EmbeddingLayer(
            latent_size=latent_size,
            num_patches=self.num_patches,
            input_dim=input_dim,
        )
        
        # Multi-Head Self-Attention
        self.multi_head_attn = nn.MultiheadAttention(
            embed_dim=latent_size,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True,  # Input: (batch, seq, feature)
        )
        
        self.norm_1 = nn.LayerNorm(latent_size)
        self.norm_2 = nn.LayerNorm(latent_size)
        self.dropout = nn.Dropout(dropout)

        # Feed Forward Block: mở rộng → thu hẹp
        self.feed_forward = nn.Sequential(
            nn.Linear(latent_size, latent_size * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(latent_size * 4, latent_size),
        )
        
        # Classification head
        self.output_layer = nn.Linear(latent_size * self.num_patches, num_classes)

    def forward(self, x):
        # Bước 1: Tách ảnh thành patches
        x = self.patchifier(x)                          # (B, N, P*P*C)
        
        # Bước 2: Embedding + Positional Encoding
        x = self.embedding_layer(x)                     # (B, N, latent_size)
        
        # Bước 3: Multi-Head Self-Attention với residual connection
        attn_out, _ = self.multi_head_attn(x, x, x)
        x = self.norm_1(self.dropout(attn_out) + x)     # Pre-norm + residual
        
        # Bước 4: Feed Forward với residual connection
        x = self.norm_2(self.feed_forward(x) + x)       # (B, N, latent_size)
        
        # Bước 5: Flatten và phân loại
        x = x.flatten(start_dim=1)                      # (B, N*latent_size)
        x = self.output_layer(x)                         # (B, num_classes)
        return x

import torchvision
import torchvision.transforms as transforms
from torch import optim
from tqdm import tqdm

# --- Data ---
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

BATCH_SIZE = 64

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
testset  = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
testloader  = torch.utils.data.DataLoader(testset,  batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

CLASSES = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# --- Model ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = ViT(
    patch_size=8,           # 32/8 = 4 patches mỗi chiều → 16 patches tổng
    img_dimension=(32, 32),
    latent_size=512,
    num_heads=4,
    num_classes=10,
).to(device)

print(f"Số tham số: {sum(p.numel() for p in model.parameters()):,}")

# --- Training ---
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
loss_fn   = nn.CrossEntropyLoss()

def train_epoch(loader):
    model.train()
    total_loss, correct, total = 0, 0, 0
    for x, y in tqdm(loader, desc="Training"):
        x, y = x.to(device), y.to(device)
        logits = model(x)
        loss = loss_fn(logits, y)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        correct += (logits.argmax(1) == y).sum().item()
        total += len(y)
    
    return total_loss / len(loader), correct / total

@torch.no_grad()
def evaluate(loader):
    model.eval()
    total_loss, correct, total = 0, 0, 0
    for x, y in tqdm(loader, desc="Evaluating"):
        x, y = x.to(device), y.to(device)
        logits = model(x)
        total_loss += loss_fn(logits, y).item()
        correct += (logits.argmax(1) == y).sum().item()
        total += len(y)
    
    return total_loss / len(loader), correct / total

def train(epochs=20):
    for epoch in range(epochs):
        train_loss, train_acc = train_epoch(trainloader)
        val_loss, val_acc     = evaluate(testloader)
        scheduler.step()
        
        print(f"Epoch {epoch+1:02d}/{epochs} | "
              f"Train Loss: {train_loss:.4f} Acc: {train_acc:.3f} | "
              f"Val Loss: {val_loss:.4f} Acc: {val_acc:.3f}")

train(20)

# Thay vì train from scratch, dùng torchvision hoặc timm
import timm
model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=10)