File size: 6,223 Bytes
6857ec0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """Encodeur image ViT from scratch (Phase 4).
Patch embedding + blocs transformer + **positional embedding appris**.
Choix : un ViT « plat » (patchs × transformer, sans CNN), comme dans « An Image is
Worth 16x16 Words », mais **entièrement réimplémenté**.
* **Patch embedding** : une convolution de noyau `patch_size` et de pas
`patch_size` découpe l'image en patchs non chevauchants et projette chacun en
`n_embd`. C'est une projection linéaire par patch, exprimée en conv pour la
simplicité.
* **Positional embedding appris** (et **pas** RoPE) : sur une grille 2D de taille
fixe et connue, apprendre un vecteur par position est simple, direct et
suffisant. RoPE 2D apporte peu ici. RoPE reste utilisé pour le texte.
* **CLS token** : un token appris additionnel résume l'image (utilisé par le
projecteur / la tête de classification). Sinon `pool='mean'`.
Justification complète : docs/architecture.md §3.
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
import torch.nn as nn
from src.model.block import TransformerBlock
@dataclass
class ViTConfig:
image_size: int = 32
patch_size: int = 8
in_channels: int = 3
n_layer: int = 4
n_head: int = 4
n_embd: int = 128
dropout: float = 0.0
bias: bool = False
ffn_mult: int = 4
pool: str = "cls" # cls | mean
use_rope: bool = False # positional embedding appris par défaut
class PatchEmbedding(nn.Module):
"""Découpe l'image en patchs et projette chaque patch en ``n_embd``."""
def __init__(self, image_size: int, patch_size: int, in_channels: int, n_embd: int):
super().__init__()
if image_size % patch_size != 0:
raise ValueError(
f"image_size ({image_size}) doit être divisible par patch_size ({patch_size})"
)
self.image_size = image_size
self.patch_size = patch_size
self.grid = image_size // patch_size
self.n_patches = self.grid * self.grid
self.proj = nn.Conv2d(in_channels, n_embd, kernel_size=patch_size, stride=patch_size)
def forward(self, images: torch.Tensor) -> torch.Tensor:
# images: (B, C, H, W) -> (B, n_patches, n_embd)
x = self.proj(images) # (B, n_embd, grid, grid)
x = x.flatten(2).transpose(1, 2) # (B, n_patches, n_embd)
return x
class VisionEncoder(nn.Module):
"""Encodeur d'images : patch embedding + position apprise + N blocs."""
def __init__(self, config: ViTConfig):
super().__init__()
self.config = config
self.patch_embed = PatchEmbedding(
config.image_size, config.patch_size, config.in_channels, config.n_embd
)
n_tokens = self.patch_embed.n_patches + (1 if config.pool == "cls" else 0)
self.cls_token = (
nn.Parameter(torch.zeros(1, 1, config.n_embd)) if config.pool == "cls" else None
)
self.pos_embedding = nn.Parameter(torch.zeros(1, n_tokens, config.n_embd))
self.dropout = nn.Dropout(config.dropout)
self.blocks = nn.ModuleList(
[
TransformerBlock(
config.n_embd,
config.n_head,
dropout=config.dropout,
bias=config.bias,
ffn_mult=config.ffn_mult,
use_rope=config.use_rope,
causal=False,
)
for _ in range(config.n_layer)
]
)
self.ln_f = nn.LayerNorm(config.n_embd, bias=config.bias)
self._init_weights()
def _init_weights(self) -> None:
nn.init.normal_(self.pos_embedding, std=0.02)
if self.cls_token is not None:
nn.init.normal_(self.cls_token, std=0.02)
for module in self.modules():
if isinstance(module, nn.Linear):
nn.init.normal_(module.weight, std=0.02)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Conv2d):
nn.init.normal_(module.weight, std=0.02)
if module.bias is not None:
nn.init.zeros_(module.bias)
def forward(self, images: torch.Tensor) -> torch.Tensor:
"""Retourne les features ``(B, n_tokens, n_embd)`` (CLS inclus si ``pool='cls'``)."""
x = self.patch_embed(images)
if self.cls_token is not None:
cls = self.cls_token.expand(x.size(0), -1, -1)
x = torch.cat([cls, x], dim=1)
x = self.dropout(x + self.pos_embedding)
for block in self.blocks:
x = block(x)
return self.ln_f(x)
def pool(self, tokens: torch.Tensor) -> torch.Tensor:
"""Réduit les tokens en un vecteur ``(B, n_embd)``."""
if self.config.pool == "cls":
return tokens[:, 0]
return tokens.mean(dim=1)
class ImageClassifier(nn.Module):
"""Tête de classification (utilisée pour l'overfit de test, Phase 4)."""
def __init__(self, config: ViTConfig, n_classes: int):
super().__init__()
self.encoder = VisionEncoder(config)
self.head = nn.Linear(config.n_embd, n_classes)
def forward(self, images: torch.Tensor) -> torch.Tensor:
tokens = self.encoder(images)
return self.head(self.encoder.pool(tokens))
class ImageDecoder(nn.Module):
"""Tête de reconstruction (auto-encodeur jouet), pour prouver que l'encodeur
capture de l'information sans utiliser d'étiquette."""
def __init__(self, config: ViTConfig, hidden: int = 256):
super().__init__()
self.encoder = VisionEncoder(config)
out_dim = config.in_channels * config.image_size * config.image_size
self.net = nn.Sequential(
nn.Linear(config.n_embd, hidden), nn.GELU(), nn.Linear(hidden, out_dim)
)
self.image_size = config.image_size
self.in_channels = config.in_channels
def forward(self, images: torch.Tensor) -> torch.Tensor:
feat = self.encoder.pool(self.encoder(images))
out = self.net(feat)
return out.view(-1, self.in_channels, self.image_size, self.image_size)
|