173 lines
7.5 KiB
Python
173 lines
7.5 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from abc import ABC, abstractmethod
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
class FrozenBackbone(nn.Module, ABC):
|
|
vision_dim: int
|
|
text_dim: int
|
|
|
|
@abstractmethod
|
|
def encode_image(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: ...
|
|
|
|
@abstractmethod
|
|
def encode_text(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]: ...
|
|
|
|
def train(self, mode: bool = True) -> "FrozenBackbone":
|
|
# The extractor trains around the backbone; pretrained backbone behavior stays fixed.
|
|
return super().train(False)
|
|
|
|
|
|
class TinyFrozenBackbone(FrozenBackbone):
|
|
"""Offline backbone used to exercise the complete pipeline in CPU CI."""
|
|
|
|
def __init__(self, hidden_dim: int = 64, patch_size: int = 16):
|
|
super().__init__()
|
|
generator = torch.Generator().manual_seed(1234)
|
|
self.patch = nn.Conv2d(3, hidden_dim, patch_size, stride=patch_size, bias=False)
|
|
self.byte_embedding = nn.Embedding(257, hidden_dim, padding_idx=0)
|
|
with torch.no_grad():
|
|
self.patch.weight.normal_(generator=generator, std=0.04)
|
|
self.byte_embedding.weight.normal_(generator=generator, std=0.04)
|
|
self.byte_embedding.weight[0].zero_()
|
|
self.requires_grad_(False)
|
|
self.vision_dim = hidden_dim
|
|
self.text_dim = hidden_dim
|
|
|
|
def encode_image(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
with torch.no_grad():
|
|
tokens = self.patch(images).flatten(2).transpose(1, 2)
|
|
mask = torch.ones(tokens.shape[:2], dtype=torch.bool, device=tokens.device)
|
|
return tokens, mask
|
|
|
|
def encode_text(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
|
|
encoded = []
|
|
for text in texts:
|
|
# Hash whole whitespace-delimited tokens so the offline representation preserves words.
|
|
ids = []
|
|
for token in text.lower().replace(".", "").split():
|
|
digest = hashlib.blake2b(token.encode("utf-8"), digest_size=2).digest()
|
|
ids.append((int.from_bytes(digest, "little") % 256) + 1)
|
|
encoded.append(ids or [1])
|
|
length = max(len(row) for row in encoded)
|
|
input_ids = torch.zeros(
|
|
len(encoded), length, dtype=torch.long, device=self.patch.weight.device
|
|
)
|
|
mask = torch.zeros_like(input_ids, dtype=torch.bool)
|
|
for index, row in enumerate(encoded):
|
|
input_ids[index, : len(row)] = torch.tensor(row, device=input_ids.device)
|
|
mask[index, : len(row)] = True
|
|
with torch.no_grad():
|
|
tokens = self.byte_embedding(input_ids)
|
|
return tokens, mask
|
|
|
|
|
|
class CLIPFrozenBackbone(FrozenBackbone):
|
|
"""Frozen CLIP token encoder for the first real GPU experiment."""
|
|
|
|
def __init__(self, model_name: str):
|
|
super().__init__()
|
|
try:
|
|
from transformers import AutoTokenizer, CLIPModel
|
|
except ImportError as error:
|
|
raise RuntimeError("Install the 'hf' extra to use the CLIP backbone.") from error
|
|
# Keep training and the post-training benchmark on the same cached artifact.
|
|
# Otherwise Hugging Face may fetch both the legacy .bin and safetensors weights.
|
|
self.model = CLIPModel.from_pretrained(model_name, use_safetensors=False)
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
self.model.requires_grad_(False).eval()
|
|
self.vision_dim = self.model.config.vision_config.hidden_size
|
|
self.text_dim = self.model.config.text_config.hidden_size
|
|
mean = torch.tensor([0.48145466, 0.4578275, 0.40821073])
|
|
std = torch.tensor([0.26862954, 0.26130258, 0.27577711])
|
|
self.register_buffer("image_mean", mean.view(1, 3, 1, 1), persistent=False)
|
|
self.register_buffer("image_std", std.view(1, 3, 1, 1), persistent=False)
|
|
|
|
def encode_image(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
pixels = (images - self.image_mean) / self.image_std
|
|
with torch.no_grad():
|
|
tokens = self.model.vision_model(pixel_values=pixels).last_hidden_state
|
|
mask = torch.ones(tokens.shape[:2], dtype=torch.bool, device=tokens.device)
|
|
return tokens, mask
|
|
|
|
def encode_text(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
|
|
values = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
|
|
values = {key: value.to(self.image_mean.device) for key, value in values.items()}
|
|
with torch.no_grad():
|
|
tokens = self.model.text_model(**values).last_hidden_state
|
|
return tokens, values["attention_mask"].bool()
|
|
|
|
|
|
class SigLIP2FrozenBackbone(FrozenBackbone):
|
|
"""Billion-scale paired image/text encoder used for formal experiments."""
|
|
|
|
def __init__(self, model_name: str):
|
|
super().__init__()
|
|
try:
|
|
from transformers import AutoModel, AutoProcessor
|
|
except ImportError as error:
|
|
raise RuntimeError("Install the 'hf' extra to use SigLIP2.") from error
|
|
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
|
|
self.model = AutoModel.from_pretrained(
|
|
model_name,
|
|
torch_dtype=dtype,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
self.processor = AutoProcessor.from_pretrained(model_name)
|
|
self.model.requires_grad_(False).eval()
|
|
self.vision_dim = self.model.config.vision_config.hidden_size
|
|
self.text_dim = self.model.config.text_config.hidden_size
|
|
|
|
def _to_model_device(self, values: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
|
|
device = next(self.model.parameters()).device
|
|
dtype = next(self.model.parameters()).dtype
|
|
moved = {}
|
|
for key, value in values.items():
|
|
moved[key] = (
|
|
value.to(device=device, dtype=dtype)
|
|
if value.is_floating_point()
|
|
else value.to(device=device)
|
|
)
|
|
return moved
|
|
|
|
def encode_image(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
values = self.processor(images=images.detach().cpu(), return_tensors="pt")
|
|
allowed = {"pixel_values", "pixel_attention_mask", "spatial_shapes"}
|
|
vision_values = self._to_model_device(
|
|
{key: value for key, value in values.items() if key in allowed}
|
|
)
|
|
with torch.no_grad():
|
|
tokens = self.model.vision_model(**vision_values).last_hidden_state
|
|
mask = torch.ones(tokens.shape[:2], dtype=torch.bool, device=tokens.device)
|
|
return tokens, mask
|
|
|
|
def encode_text(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
|
|
max_length = self.model.config.text_config.max_position_embeddings
|
|
values = self.processor(
|
|
text=texts,
|
|
padding="max_length",
|
|
truncation=True,
|
|
max_length=max_length,
|
|
return_tensors="pt",
|
|
)
|
|
text_values = self._to_model_device(
|
|
{key: value for key, value in values.items() if key in {"input_ids", "attention_mask"}}
|
|
)
|
|
with torch.no_grad():
|
|
tokens = self.model.text_model(**text_values).last_hidden_state
|
|
return tokens, text_values["attention_mask"].bool()
|
|
|
|
|
|
def build_backbone(name: str, model_name: str) -> FrozenBackbone:
|
|
if name == "tiny":
|
|
return TinyFrozenBackbone()
|
|
if name == "clip":
|
|
return CLIPFrozenBackbone(model_name)
|
|
if name == "siglip2":
|
|
return SigLIP2FrozenBackbone(model_name)
|
|
raise ValueError(f"Unknown backbone: {name}")
|