107 lines
4.6 KiB
Python
107 lines
4.6 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
|
|
self.model = CLIPModel.from_pretrained(model_name)
|
|
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()
|
|
|
|
|
|
def build_backbone(name: str, clip_model: str) -> FrozenBackbone:
|
|
if name == "tiny":
|
|
return TinyFrozenBackbone()
|
|
if name == "clip":
|
|
return CLIPFrozenBackbone(clip_model)
|
|
raise ValueError(f"Unknown backbone: {name}")
|