Spaces:
Sleeping
Sleeping
File size: 1,775 Bytes
5e94db5 70e1e32 5e94db5 70e1e32 5e94db5 |
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 |
from __future__ import annotations
import numpy as np
import torch
import open_clip
from contextlib import nullcontext
from src.models.utils import l2norm_rows
class CLIPLinearProbe:
def __init__(self, head_path):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
self.model, _, self.preprocess = open_clip.create_model_and_transforms(
"ViT-L-14", pretrained="openai", device=self.device
)
self.model.eval().requires_grad_(False)
npz = np.load(head_path)
self.w = torch.from_numpy(npz["w"]).to(self.device).float()
self.b = torch.from_numpy(npz["b"]).to(self.device).float()
self.use_amp = False
if self.device == "cuda":
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.benchmark = True
self.use_amp = True
@torch.inference_mode()
def encode(self, pil_list) -> torch.Tensor:
x = torch.stack([self.preprocess(im.convert("RGB")) for im in pil_list], 0)
x = x.to(self.device, non_blocking=True, memory_format=torch.channels_last)
ctx = torch.amp.autocast("cuda", dtype=self.torch_dtype) if self.use_amp else nullcontext()
with ctx:
f = self.model.encode_image(x)
f = f.float()
return l2norm_rows(f)
@torch.inference_mode()
def logits(self, pil_list) -> torch.Tensor:
f = self.encode(pil_list)
return (f @ self.w + self.b).squeeze(1)
@torch.inference_mode()
def prob(self, pil_list) -> torch.Tensor:
z = torch.clamp(self.logits(pil_list), -50, 50)
return torch.sigmoid(z)
|