Spaces:
Sleeping
Sleeping
File size: 5,151 Bytes
f343f06 | 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | """
Embedding cache utilities for pre-computed features.
Handles saving/loading embeddings to disk for fast retrieval.
"""
import torch
from pathlib import Path
from typing import Tuple, List, Optional, Dict
import json
def save_embeddings(
embeddings: torch.Tensor,
metadata: Dict,
output_dir: str,
filename: Optional[str] = None
) -> Path:
"""
Save embeddings and metadata to disk.
Args:
embeddings: Tensor of shape (N, embed_dim)
metadata: Dict with keys like 'modality', 'sample_ids', 'class_labels'
output_dir: Directory to save to
filename: Optional custom filename (without extension)
Returns:
Path to saved file
"""
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
# Generate filename if not provided
if filename is None:
modality = metadata.get("modality", "unknown")
n_samples = embeddings.shape[0]
filename = f"{modality}_embeddings_{n_samples}"
# Save embeddings tensor
embeddings_path = output_dir / f"{filename}.pt"
torch.save(embeddings, embeddings_path)
# Save metadata as JSON
metadata_path = output_dir / f"{filename}_metadata.json"
# Convert tensors in metadata to lists for JSON serialization
serializable_metadata = {}
for key, value in metadata.items():
if isinstance(value, torch.Tensor):
serializable_metadata[key] = value.tolist()
else:
serializable_metadata[key] = value
with open(metadata_path, "w") as f:
json.dump(serializable_metadata, f, indent=2)
return embeddings_path
def load_embeddings(
cache_path: str
) -> Tuple[torch.Tensor, Dict]:
"""
Load embeddings and metadata from disk.
Args:
cache_path: Path to .pt embeddings file
Returns:
(embeddings tensor, metadata dict)
"""
cache_path = Path(cache_path)
# Load embeddings
embeddings = torch.load(cache_path, weights_only=True)
# Load metadata if exists
metadata_path = cache_path.with_name(
cache_path.stem + "_metadata.json"
)
metadata = {}
if metadata_path.exists():
with open(metadata_path, "r") as f:
metadata = json.load(f)
return embeddings, metadata
def get_cache_path(
output_dir: str,
modality: str,
split: str = "gallery",
embed_dim: int = 768
) -> Path:
"""
Generate standard cache file path.
Args:
output_dir: Base output directory
modality: Modality type
split: Dataset split (query/gallery)
embed_dim: Embedding dimension
Returns:
Path object for cache file
"""
output_dir = Path(output_dir)
filename = f"{modality}_{split}_embeddings.pt"
return output_dir / filename
def verify_embeddings(
embeddings: torch.Tensor,
expected_dim: Optional[int] = None,
l2_normalized: bool = True
) -> bool:
"""
Verify embeddings are valid.
Args:
embeddings: Embedding tensor
expected_dim: Expected embedding dimension
l2_normalized: Whether embeddings should be L2-normalized
Returns:
True if valid
"""
if embeddings.dim() != 2:
print(f"Expected 2D tensor, got {embeddings.dim()}D")
return False
if expected_dim is not None and embeddings.shape[1] != expected_dim:
print(f"Expected dim {expected_dim}, got {embeddings.shape[1]}")
return False
if l2_normalized:
norms = torch.norm(embeddings, dim=1)
# Check if norms are close to 1 (allowing for floating point)
if not torch.allclose(norms, torch.ones_like(norms), atol=1e-3):
print(f"Embeddings not L2-normalized. Norms: {norms[:5]}")
return False
return True
# Self-check
if __name__ == "__main__":
import tempfile
print("Testing embedding cache utilities...")
# Create dummy embeddings
embeddings = torch.randn(100, 768)
embeddings = torch.nn.functional.normalize(embeddings, dim=1) # L2 normalize
metadata = {
"modality": "optical",
"sample_ids": list(range(100)),
"class_labels": [i % 10 for i in range(100)],
}
# Test save/load roundtrip
with tempfile.TemporaryDirectory() as tmpdir:
# Save
save_path = save_embeddings(embeddings, metadata, tmpdir, "test")
print(f"Saved to: {save_path}")
# Load
loaded_embeddings, loaded_metadata = load_embeddings(save_path)
print(f"Loaded embeddings shape: {loaded_embeddings.shape}")
print(f"Loaded metadata keys: {list(loaded_metadata.keys())}")
# Verify
assert torch.allclose(embeddings, loaded_embeddings), "Embeddings mismatch!"
assert loaded_metadata["modality"] == "optical", "Metadata mismatch!"
# Test verify
assert verify_embeddings(loaded_embeddings, expected_dim=768), "Verification failed!"
print("\nEmbedding cache test passed!")
|