File size: 4,233 Bytes
dac64fd a403c2d dac64fd a403c2d dac64fd a403c2d dac64fd | 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 | """Standalone loader for clip-vitb-mini-distilled (no repo imports).
Usage:
from loader import load_model, embed_images
model = load_model("student.pt") # or a hf_hub_download path
emb = embed_images(model, batch01) # (B,3,H,W) in [0,1] -> (B,512)
The student is an 8.66M-parameter ViT (d=240, depth 12, heads 4, patch 16,
img 160) distilled from CLIP ViT-B/16 (LAION-2B) pooled image features on
COCO-2017 train (118,287 images, ~17 epochs). Outputs live in the teacher's
512-d projection space: compatible with the teacher's TEXT tower for
zero-shot/retrieval. See the model card for measured capability and caveats.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
CLIP_STD = (0.26862954, 0.26130258, 0.27577711)
class Block(nn.Module):
def __init__(self, d, heads):
super().__init__()
self.n1 = nn.LayerNorm(d)
self.qkv = nn.Linear(d, 3 * d)
self.proj = nn.Linear(d, d)
self.n2 = nn.LayerNorm(d)
self.fc1 = nn.Linear(d, 4 * d)
self.fc2 = nn.Linear(4 * d, d)
self.heads = heads
def forward(self, x):
B, N, C = x.shape
q, k, v = (self.qkv(self.n1(x))
.reshape(B, N, 3, self.heads, C // self.heads)
.permute(2, 0, 3, 1, 4))
a = F.scaled_dot_product_attention(q, k, v)
x = x + self.proj(a.transpose(1, 2).reshape(B, N, C))
return x + self.fc2(F.gelu(self.fc1(self.n2(x))))
class Student(nn.Module):
"""CLS-token readout, Linear head to the teacher's 512-d space."""
def __init__(self, out_dim=512, d=240, depth=12, heads=4, patch=16,
img=160):
super().__init__()
self.patch = nn.Conv2d(3, d, patch, patch)
self.cls = nn.Parameter(torch.zeros(1, 1, d))
self.pos = nn.Parameter(torch.zeros(1, (img // patch) ** 2 + 1, d))
self.blocks = nn.ModuleList(Block(d, heads) for _ in range(depth))
self.norm = nn.LayerNorm(d)
self.head = nn.Linear(d, out_dim)
def forward_features(self, x):
x = self.patch(x).flatten(2).transpose(1, 2)
x = torch.cat([self.cls.expand(x.shape[0], -1, -1), x], 1) + self.pos
for b in self.blocks:
x = b(x)
return self.norm(x)[:, 0]
def forward(self, x):
return self.head(self.forward_features(x))
def load_model(path, device=None):
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
ck = torch.load(path, map_location="cpu", weights_only=True)
sd = ck["state_dict"] if "state_dict" in ck else ck
model = Student(out_dim=sd["head.weight"].shape[0])
model.load_state_dict(sd, strict=True)
return model.to(device).eval()
def load_rotation(path):
"""Load a frame rotation shipped beside a consensus-distilled student.
The rotation maps the student's consensus frame into the CLIP-B/16
(LAION-2B) deployment frame; it was fitted once, offline, on 2,500
COCO-val pairs (fp64 orthogonal Procrustes) and is FROZEN."""
d = torch.load(path, map_location="cpu", weights_only=True)
return d["R"] if isinstance(d, dict) else d
@torch.no_grad()
def embed_images(model, x01, batch=256, rotation=None):
"""x01: (B,3,H,W) float in [0,1]. Returns L2-normalized (B,512).
Pass `rotation` (from load_rotation) with the consensus-distilled
champion to place outputs in the LAION text tower's frame — required
for zero-shot/retrieval against that tower, and how every headline
number on the card was measured."""
dev = next(model.parameters()).device
mean = torch.tensor(CLIP_MEAN).view(1, 3, 1, 1).to(dev)
std = torch.tensor(CLIP_STD).view(1, 3, 1, 1).to(dev)
out = []
for i in range(0, len(x01), batch):
x = x01[i:i + batch].to(dev)
x = F.interpolate(x, size=(160, 160), mode="bicubic",
align_corners=False).clamp(0, 1)
z = F.normalize(model((x - mean) / std), dim=-1)
if rotation is not None:
z = F.normalize(z.double() @ rotation.to(dev).double(),
dim=-1).float()
out.append(z.cpu())
return torch.cat(out)
|