File size: 5,363 Bytes
827713a | 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 | """Modèle CLIP maison (Phase 5) : alignement image-texte from scratch.
Deux encodeurs (texte + image) produisent chacun un vecteur, projeté dans un espace
commun puis **normalisé L2**. La similarité est un produit scalaire. Aucun modèle
CLIP pré-entraîné n'est utilisé : encodeur image = `VisionEncoder` maison (Phase 4),
encodeur texte = transformer maison (Phase 2).
Justification : docs/architecture.md §4.
"""
from __future__ import annotations
from dataclasses import dataclass, field
import torch
import torch.nn as nn
import torch.nn.functional as F
from src.model.projectors import MLPProjector
from src.model.text_model import CortexTextLM, TextConfig
from src.model.vit import VisionEncoder, ViTConfig
def _default_text_config() -> TextConfig:
return TextConfig(vocab_size=32000, n_layer=2, n_head=4, n_embd=128,
ctx_len=32, tie_embeddings=True)
def _default_vision_config() -> ViTConfig:
return ViTConfig(image_size=32, patch_size=8, n_layer=2, n_head=4, n_embd=128)
@dataclass
class ClipConfig:
embed_dim: int = 64 # dimension de l'espace commun (image/texte)
text: TextConfig = field(default_factory=_default_text_config)
vision: ViTConfig = field(default_factory=_default_vision_config)
init_temperature: float = 0.07
learn_temperature: bool = True
class TextEncoder(nn.Module):
"""Transformer texte (Phase 2) -> pooling masqué -> vecteur ``(B, n_embd)``."""
def __init__(self, config: TextConfig):
super().__init__()
self.backbone = CortexTextLM(config)
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None):
hidden = self.backbone.encode(input_ids) # (B, T, n_embd), avant la tête
if attention_mask is None:
return hidden.mean(dim=1)
mask = attention_mask.unsqueeze(-1).float()
return (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-6)
class CortexClip(nn.Module):
"""CLIP maison : encodeurs texte + image, projecteurs, `logit_scale` apprise."""
def __init__(self, config: ClipConfig):
super().__init__()
self.config = config
self.text_encoder = TextEncoder(config.text)
self.vision_encoder = VisionEncoder(config.vision)
self.text_proj = MLPProjector(config.text.n_embd, config.embed_dim)
self.image_proj = MLPProjector(config.vision.n_embd, config.embed_dim)
init_logit_scale = float(torch.log(torch.tensor(1.0 / config.init_temperature)))
if config.learn_temperature:
self.logit_scale = nn.Parameter(torch.tensor(init_logit_scale))
else:
self.register_buffer("logit_scale", torch.tensor(init_logit_scale))
def encode_text(self, input_ids, attention_mask=None):
pooled = self.text_proj(self.text_encoder(input_ids, attention_mask))
return F.normalize(pooled, dim=-1)
def encode_image(self, images):
feat = self.vision_encoder.pool(self.vision_encoder(images))
return F.normalize(self.image_proj(feat), dim=-1)
def forward(self, images, input_ids, attention_mask=None):
img = self.encode_image(images)
txt = self.encode_text(input_ids, attention_mask)
scale = self.logit_scale.exp().clamp(max=100.0)
logits_per_image = scale * img @ txt.t()
return logits_per_image, logits_per_image.t(), img, txt
def contrastive_loss(logits_per_image: torch.Tensor, logits_per_text: torch.Tensor) -> torch.Tensor:
"""Perte contrastive symétrique (écrite à la main).
La bonne paire est sur la diagonale : image i ↔ texte i.
"""
batch = logits_per_image.size(0)
targets = torch.arange(batch, device=logits_per_image.device)
return 0.5 * (F.cross_entropy(logits_per_image, targets)
+ F.cross_entropy(logits_per_text, targets))
def clip_loss(model: CortexClip, images, input_ids, attention_mask=None) -> torch.Tensor:
logits_per_image, logits_per_text, _, _ = model(images, input_ids, attention_mask)
return contrastive_loss(logits_per_image, logits_per_text)
@torch.no_grad()
def retrieval_accuracy(model: CortexClip, images, input_ids, attention_mask=None) -> float:
"""Précision de retrieval image->texte (diagonale) sur un lot **sans doublon**."""
img = model.encode_image(images)
txt = model.encode_text(input_ids, attention_mask)
preds = (img @ txt.t()).argmax(dim=-1)
targets = torch.arange(images.size(0), device=images.device)
return (preds == targets).float().mean().item()
@torch.no_grad()
def class_retrieval_accuracy(
model: CortexClip,
images: torch.Tensor,
bank_ids: torch.Tensor,
bank_mask: torch.Tensor,
label_ids: torch.Tensor,
) -> float:
"""Retrieval image->texte **au niveau des classes**.
``bank_ids``/``bank_mask`` décrivent la banque de légendes **uniques** (ici les
9 classes) ; ``label_ids[i]`` est l'indice de la vraie légende de l'image ``i``
dans cette banque. C'est la bonne mesure quand les légendes se répètent (une
légende partagée par 10 images ne peut pas être « la » bonne par un argmax
sur 90 colonnes).
"""
img = model.encode_image(images)
txt = model.encode_text(bank_ids, bank_mask)
preds = (img @ txt.t()).argmax(dim=-1)
return (preds == label_ids).float().mean().item()
|