File size: 5,469 Bytes
534cc76 | 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 | """Cache frozen FP32 features by signal, preprocessing and encoder identity."""
from pathlib import Path
import hashlib
import json
import sqlite3
import torch
from .provenance import CACHE_CONTRACT, fingerprint
from .atomic import atomic_output
def tensor_digest(values):
import hashlib
digest = hashlib.sha256()
digest.update(values.contiguous().numpy().tobytes())
return digest.hexdigest()
class FeatureCache:
def __init__(self, directory):
self.directory = Path(directory)
self.directory.mkdir(parents=True, exist_ok=True)
def get_or_encode(self, model, prepared):
binding = {"contract": CACHE_CONTRACT, "model": model.identity(),
"prepared": prepared.identity}
path = self.directory / (fingerprint(binding) + ".pt")
if path.exists():
payload = torch.load(path, map_location="cpu", weights_only=True)
if payload.get("binding") != binding:
raise ValueError("Frozen feature cache contract mismatch")
values = payload["values"]
if values.shape != (1, 512):
raise ValueError("Frozen feature cache shape mismatch")
if values.dtype != torch.float32:
raise ValueError("Frozen feature cache must preserve FP32 computation")
if not torch.isfinite(values).all():
raise ValueError("Nonfinite cached features")
if tensor_digest(values) != payload.get("digest"):
raise ValueError("Cached feature content changed")
return values.to(model.device)
encoded = model.encode(prepared)
values = encoded.cpu().float()
with atomic_output(path) as output:
torch.save({"binding": binding, "values": values,
"digest": tensor_digest(values)}, output)
return encoded
class SQLiteFeatureCache:
"""Batch native encodes and keep verified FP32 features in one durable store."""
def __init__(self, path):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
self.db = sqlite3.connect(path)
self.db.executescript('''PRAGMA journal_mode=WAL; PRAGMA synchronous=FULL;
PRAGMA cache_size=-16384; PRAGMA temp_store=FILE;
CREATE TABLE IF NOT EXISTS metadata(binding TEXT PRIMARY KEY);
CREATE TABLE IF NOT EXISTS features(id TEXT PRIMARY KEY, values_blob BLOB, digest TEXT);''')
@torch.no_grad()
def get_many_or_encode(self, model, prepared):
identity = dict(model.identity())
identity.pop('head', None)
current = {"contract":CACHE_CONTRACT,"model":identity}
binding = json.dumps(current,sort_keys=True)
existing = self.db.execute('SELECT binding FROM metadata').fetchone()
if existing:
previous = json.loads(existing[0])
# Old stores bound head settings even though frozen features do not depend on them.
if isinstance(previous.get('model'),dict):
previous['model'].pop('head',None)
if previous != current:
raise ValueError("Frozen feature cache contract mismatch")
if not existing:
with self.db: self.db.execute('INSERT INTO metadata VALUES (?)',(binding,))
unique = {value.identity:value for value in prepared}
found = {}
keys = list(unique)
for start in range(0,len(keys),500):
chunk=keys[start:start+500]
for key,content,digest in self.db.execute(
'SELECT id,values_blob,digest FROM features WHERE id IN ('+','.join('?' for _ in chunk)+')',chunk):
if len(content)!=512*4 or hashlib.sha256(content).hexdigest()!=digest:
raise ValueError("Cached feature content changed")
values=torch.frombuffer(bytearray(content),dtype=torch.float32).clone()
if not torch.isfinite(values).all():
raise ValueError("Nonfinite cached features")
found[key]=values
shapes = {}
for key,value in unique.items():
if key not in found:
shape=tuple(model.preprocessor.format(value).samples.shape)
shapes.setdefault(shape,[]).append((key,value))
for entries in shapes.values():
encoded=model.encode_batch([value for _,value in entries]).cpu().float()
updates=[]
for (key,_),values in zip(entries,encoded):
content=values.contiguous().numpy().tobytes()
updates.append((key,content,hashlib.sha256(content).hexdigest()))
found[key]=values
with self.db:
self.db.executemany('INSERT INTO features VALUES (?,?,?)',updates)
return torch.stack([found[value.identity] for value in prepared]).to(model.device)
def get_or_encode(self, model, prepared):
return self.get_many_or_encode(model,[prepared])
def open_feature_cache(config):
if not config.get("feature_cache"):
return None
kind=config.get("feature_cache_format","files")
if kind=="sqlite": return SQLiteFeatureCache(config["feature_cache"])
if kind=="files": return FeatureCache(config["feature_cache"])
raise ValueError("feature_cache_format must be files or sqlite")
|