Vision Transformers

import torch
import torch.nn as nn
import math

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_ch=3, embed_dim=768):
        super().__init__()
        self.n_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_ch, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        x = self.proj(x).flatten(2).transpose(1, 2)  # (B, N, D)
        return x

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_dim, n_heads, dropout=0.0):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = embed_dim // n_heads
        self.scale = self.head_dim ** -0.5
        self.qkv = nn.Linear(embed_dim, 3 * embed_dim)
        self.proj = nn.Linear(embed_dim, embed_dim)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        B, N, D = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.n_heads, self.head_dim).permute(2,0,3,1,4)
        q, k, v = qkv.unbind(0)
        attn = (q @ k.transpose(-2,-1)) * self.scale
        attn = self.drop(attn.softmax(dim=-1))
        return (attn @ v).transpose(1,2).reshape(B, N, D)

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, n_heads, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn  = MultiHeadSelfAttention(embed_dim, n_heads)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.mlp   = nn.Sequential(
            nn.Linear(embed_dim, int(embed_dim*mlp_ratio)), nn.GELU(),
            nn.Linear(int(embed_dim*mlp_ratio), embed_dim),
        )
    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

class VisionTransformer(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_ch=3, num_classes=1000,
                 embed_dim=768, depth=12, n_heads=12):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size, in_ch, embed_dim)
        N = self.patch_embed.n_patches
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, N+1, embed_dim))
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        self.blocks = nn.ModuleList([TransformerBlock(embed_dim, n_heads) for _ in range(depth)])
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)

    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)
        x = torch.cat([self.cls_token.expand(B,-1,-1), x], dim=1) + self.pos_embed
        for block in self.blocks: x = block(x)
        return self.head(self.norm(x[:, 0]))

# ViT-B/16: 224x224 → 196 patches
vit = VisionTransformer()
print(vit(torch.randn(2, 3, 224, 224)).shape)  # (2, 1000)

import torch
import torch.nn.functional as F
from torchvision import transforms

# DINOv2 feature extraction
dinov2 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
dinov2.eval()

transform = transforms.Compose([
    transforms.Resize(224), transforms.CenterCrop(224), transforms.ToTensor(),
    transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),
])

from PIL import Image
img = transform(Image.open("photo.jpg").convert("RGB")).unsqueeze(0)
with torch.no_grad():
    feats = dinov2.forward_features(img)
cls_feat   = feats["x_norm_clstoken"]    # (1, 384) for classification
patch_feat = feats["x_norm_patchtokens"] # (1, 256, 384) for dense tasks

# ViT fine-tuning with HuggingFace
from transformers import ViTForImageClassification, ViTImageProcessor, TrainingArguments, Trainer
from datasets import load_dataset
from sklearn.metrics import accuracy_score

dataset = load_dataset("food101", split={"train": "train[:5000]", "test": "validation[:500]"})
processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")

def preprocess(batch):
    inputs = processor(images=[img.convert("RGB") for img in batch["image"]], return_tensors="pt")
    inputs["labels"] = batch["label"]
    return inputs

dataset = dataset.map(preprocess, batched=True, remove_columns=["image"])
dataset.set_format("torch")

model = ViTForImageClassification.from_pretrained(
    "google/vit-base-patch16-224",
    num_labels=101, ignore_mismatched_sizes=True,
)
args = TrainingArguments(
    output_dir="vit-food101", per_device_train_batch_size=32,
    num_train_epochs=5, learning_rate=2e-5, evaluation_strategy="epoch",
)
trainer = Trainer(
    model=model, args=args,
    train_dataset=dataset["train"], eval_dataset=dataset["test"],
    compute_metrics=lambda ep: {"accuracy": accuracy_score(ep.label_ids, ep.predictions.argmax(-1))},
)
trainer.train()