diff --git a/CLAUDE.md b/CLAUDE.md index 4be2a94..f244857 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -101,7 +101,9 @@ Eval: Resize(256) + CenterCrop(256) + ImageNet normalization. | `src/models/asymmetric_encoder.py` | DINOv3ViT + TextFusionMLP + AsymmetricEncoder + GatedFusion + encode_query/encode_gallery (per-sample caption mask) | | `src/datasets/gtauav_dataset.py` | GTA-UAV-LR loader + L1/L2/L3 captions + GTAUAVSatGallery/GTAUAVDroneQuery (full retrieval eval) | | `src/datasets/mutually_exclusive_sampler.py` | BatchSampler: drone'ы в батче не делят sat_candidates (no false negatives) | -| `src/datasets/dynamic_similarity_sampler.py` | DSS: embedding-kNN + mutex — батчи из визуально похожих drone'ов (hard negatives) | +| `src/datasets/dynamic_similarity_sampler.py` | DSS: embedding-kNN + mutex — батчи из визуально похожих drone'ов (GPU/CPU, опциональный LSH) | +| `src/datasets/lsh_index.py` | Random-projection cosine-LSH для approximate kNN (opt-in; `dss_use_lsh=True`) | +| `src/datasets/embedding_cache.py` | Дисковый кеш для drone embeddings — skip re-embed на resume | | `src/losses/multi_infonce.py` | **Primary:** SymmetricInfoNCE + MoCo queue, learnable τ clamp [0.01, 0.1], weights q2g=0.6 g2q=0.4, `hard_mining_k` для top-K hardest negatives | | `src/losses/weighted_infonce.py` | Alternative: per-sample adaptive label smoothing (активируется `loss_type="weighted"`) | | `src/losses/hard_negatives.py` | NegativeMemoryBank (MoCo-style FIFO queue 4096 × 1024) | @@ -214,7 +216,7 @@ Meta-файл `meta/seg_filter.json`: исключение изображени - **Scheduler:** linear warmup (2 epochs) + cosine annealing (per-step) - **Loss:** SymmetricInfoNCE (q2g=0.6, g2q=0.4) с learnable τ (init=0.07, clamp [0.01, 0.1]) - **Hard mining:** top-K=512 hardest negatives per query из MoCo queue (размер 4096); `hard_mining_k=0` отключает -- **Batch sampler:** `sampler_type="dss"` (default) — DynamicSimilaritySampler с re-embedding каждую эпоху: пакует визуально похожих drone'ов в один батч (+hardness) с mutex-constraint (no false negatives). Первая эпоха warmup mutex-only. Средний in-batch cosine sim ~0.71 vs 0.26 у mutex. `sampler_type="mutex"` — без DSS +- **Batch sampler:** `sampler_type="dss"` (default) — DynamicSimilaritySampler с re-embedding каждую эпоху: пакует визуально похожих drone'ов в один батч (+hardness) с mutex-constraint (no false negatives). Первая эпоха warmup mutex-only. Средний in-batch cosine sim ~0.71 vs 0.26 у mutex. kNN на GPU (`dss_knn_device="cuda"`) — 1.6s vs 17s на CPU. Опциональный LSH (`dss_use_lsh=True`) для scale 100K+. Embedding cache (`dss_cache_dir`) — skip re-embed на resume. - **Eval:** full satellite gallery (~2684 unique tiles для test_20) с multi-match R@K (учитывает все positive/semi-positive) - **Augmentations:** - Drone: RandomResizedCrop(0.7-1.0), HFlip, Rotation(15°), ColorJitter, Grayscale(5%), GaussianBlur diff --git a/src/datasets/dynamic_similarity_sampler.py b/src/datasets/dynamic_similarity_sampler.py index bb32258..e0de747 100644 --- a/src/datasets/dynamic_similarity_sampler.py +++ b/src/datasets/dynamic_similarity_sampler.py @@ -27,6 +27,8 @@ import torch import torch.nn.functional as F from torch.utils.data.sampler import Sampler +from src.datasets.lsh_index import LSHIndex + LOGGER = logging.getLogger("caption_test.dss_sampler") @@ -49,6 +51,10 @@ class DynamicSimilaritySampler(Sampler[list[int]]): shuffle: bool = True, seed: int = 0, allow_partial: bool = False, + knn_device: str = "cpu", + use_lsh: bool = False, + lsh_num_tables: int = 8, + lsh_num_bits: int = 14, ) -> None: if batch_size <= 0: raise ValueError(f"batch_size must be positive, got {batch_size}") @@ -59,8 +65,13 @@ class DynamicSimilaritySampler(Sampler[list[int]]): self.shuffle = shuffle self.seed = seed self.allow_partial = allow_partial + self.knn_device = knn_device + self.use_lsh = use_lsh + self.lsh_num_tables = lsh_num_tables + self.lsh_num_bits = lsh_num_bits self.epoch = 0 self._embeddings: torch.Tensor | None = None + self._lsh: LSHIndex | None = None self._cached_len: int | None = None def set_epoch(self, epoch: int) -> None: @@ -71,22 +82,42 @@ class DynamicSimilaritySampler(Sampler[list[int]]): """Store query embeddings [N, D] for the next epoch's sampling. Call from the training loop before starting each epoch where you want - DSS to be active. Stored on CPU as float32 and L2-normalized. + DSS to be active. Embeddings are L2-normalized and moved to + `self.knn_device` for the similarity step. If `use_lsh=True`, builds + an LSH index over the embeddings for approximate candidate lookup. """ if embeddings.size(0) != len(self._item_sats): raise ValueError( f"embeddings have {embeddings.size(0)} rows but sampler tracks " f"{len(self._item_sats)} items", ) - self._embeddings = F.normalize(embeddings.detach().cpu().float(), dim=-1) + emb = F.normalize(embeddings.detach().float(), dim=-1) + self._embeddings = emb.to(self.knn_device) self._cached_len = None + + if self.use_lsh: + # LSH builds on CPU (bucket dicts are Python-native). + self._lsh = LSHIndex( + dim=emb.size(1), + num_tables=self.lsh_num_tables, + num_bits=self.lsh_num_bits, + seed=self.seed, + ) + self._lsh.build(emb.cpu()) + else: + self._lsh = None + LOGGER.info( - "DSS embeddings updated: %d × %d", *self._embeddings.shape, + "DSS embeddings updated: %d × %d on %s%s", + self._embeddings.shape[0], self._embeddings.shape[1], + self.knn_device, + f" + LSH (tables={self.lsh_num_tables}, bits={self.lsh_num_bits})" if self.use_lsh else "", ) def clear_embeddings(self) -> None: """Revert to mutex-only sampling for the next epoch.""" self._embeddings = None + self._lsh = None self._cached_len = None def _generate_batches_mutex_only(self) -> list[list[int]]: @@ -121,10 +152,16 @@ class DynamicSimilaritySampler(Sampler[list[int]]): return batches def _generate_batches_dss(self) -> list[list[int]]: - """Similarity-guided batches using stored embeddings.""" + """Similarity-guided batches using stored embeddings. + + For each seed we compute (or retrieve from LSH) an order of + candidates ranked by cosine similarity, then greedy-pack respecting + mutex constraints. Runs on `knn_device` — GPU is strongly preferred + for large N. + """ assert self._embeddings is not None rng = random.Random(self.seed + self.epoch) if self.shuffle else None - emb = self._embeddings # already L2-normalized, on CPU + emb = self._embeddings # already L2-normalized, on knn_device n = emb.size(0) remaining = set(range(n)) @@ -137,9 +174,24 @@ class DynamicSimilaritySampler(Sampler[list[int]]): else next(iter(remaining)) ) - # Cosine similarity from seed to all items (fp32 matmul, ~25K ops). - sims = emb @ emb[seed_idx] # [N] - order = sims.argsort(descending=True).tolist() + if self._lsh is not None: + # LSH path: restrict ranking to the candidate pool for this seed + # (union of buckets across tables), then exact-rank the candidates. + cand = self._lsh.candidates(emb[seed_idx].cpu()) + cand &= remaining + if len(cand) < self.batch_size: + # Candidate pool too small — fall back to full ranking to + # preserve batch-fill guarantees. + sims = emb @ emb[seed_idx] + order = sims.argsort(descending=True).cpu().tolist() + else: + cand_list = list(cand) + sims = emb[cand_list] @ emb[seed_idx] # [|cand|] + inner_order = sims.argsort(descending=True).cpu().tolist() + order = [cand_list[i] for i in inner_order] + else: + sims = emb @ emb[seed_idx] # [N] + order = sims.argsort(descending=True).cpu().tolist() batch: list[int] = [] claimed: set[str] = set() @@ -155,8 +207,8 @@ class DynamicSimilaritySampler(Sampler[list[int]]): break if len(batch) < self.batch_size: - # Couldn't fill from this seed's neighborhood (heavy mutex - # conflict). Drop the seed and retry with remaining pool. + # Couldn't fill from this seed's neighborhood. Drop the seed + # and retry with the remaining pool. remaining.discard(seed_idx) continue @@ -165,7 +217,6 @@ class DynamicSimilaritySampler(Sampler[list[int]]): remaining.discard(i) if self.allow_partial and remaining: - # Mutex-pack whatever's left. claimed = set() tail: list[int] = [] for idx in remaining: diff --git a/src/datasets/embedding_cache.py b/src/datasets/embedding_cache.py new file mode 100644 index 0000000..b4847ea --- /dev/null +++ b/src/datasets/embedding_cache.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +"""Disk cache for drone-query embeddings used by DynamicSimilaritySampler. + +Caches the [N, D] tensor produced by `_embed_drone_queries` keyed by epoch +number, so `--resume` doesn't need to recompute and ablations that rerun +from a snapshot are reproducible. + +Entries are plain `torch.save` files; there's no atomicity beyond what the +OS provides. Safe enough for single-writer usage within the training loop. +""" + +import logging +from pathlib import Path + +import torch + +LOGGER = logging.getLogger("caption_test.emb_cache") + + +class EmbeddingCache: + """File-based cache of epoch → [N, D] CPU tensor.""" + + def __init__(self, cache_dir: str | Path) -> None: + self.cache_dir = Path(cache_dir) + self.cache_dir.mkdir(parents=True, exist_ok=True) + + def path_for(self, epoch: int) -> Path: + return self.cache_dir / f"queries_epoch{epoch:03d}.pt" + + def load(self, epoch: int) -> torch.Tensor | None: + """Return cached embeddings for `epoch` or None if absent.""" + p = self.path_for(epoch) + if not p.exists(): + return None + try: + emb = torch.load(p, map_location="cpu", weights_only=False) + except Exception as exc: + LOGGER.warning("Failed to load embedding cache %s: %s", p, exc) + return None + LOGGER.info("Loaded cached embeddings from %s (shape=%s)", p.name, tuple(emb.shape)) + return emb + + def save(self, epoch: int, embeddings: torch.Tensor) -> None: + p = self.path_for(epoch) + tmp = p.with_suffix(p.suffix + ".tmp") + torch.save(embeddings.detach().cpu(), tmp) + tmp.replace(p) + LOGGER.info("Saved embedding cache to %s (shape=%s)", p.name, tuple(embeddings.shape)) diff --git a/src/datasets/lsh_index.py b/src/datasets/lsh_index.py new file mode 100644 index 0000000..cb11ed1 --- /dev/null +++ b/src/datasets/lsh_index.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +"""Random-projection LSH for approximate cosine-similarity kNN. + +Builds H independent hash tables, each with B bits. Every embedding hashes to +a B-bit code per table based on its sign against B random hyperplanes. +Two vectors collide in a table with probability (1 − θ/π)^B where θ is the +angle between them, so cosine-similar vectors collide frequently across +tables. + +For each query we union the candidate sets from all H tables, then refine by +exact cosine similarity. Typical regime: `num_tables=8, num_bits=14` gives +recall@10 ≥ 0.95 on CVGL-scale features while keeping candidate pools under +a few percent of N. + +Not used by default at the current 25K-item scale (exact kNN is fast enough); +this module exists as a drop-in acceleration path when the dataset grows. +""" + +import logging + +import torch +import torch.nn.functional as F + +LOGGER = logging.getLogger("caption_test.lsh") + + +class LSHIndex: + """Random-projection cosine-LSH index. + + Args: + dim: Embedding dimension. + num_tables: Number of independent hash tables (more → higher recall, + larger candidate pool). + num_bits: Bits per code (more → lower collision rate per table, so + you need more tables to compensate). Typical 12-16. + seed: RNG seed for hyperplane draw (makes index reproducible). + """ + + def __init__( + self, + dim: int, + num_tables: int = 8, + num_bits: int = 14, + seed: int = 0, + ) -> None: + if num_bits > 62: + raise ValueError("num_bits must be <= 62 to fit a signed int64 hash") + self.dim = dim + self.num_tables = num_tables + self.num_bits = num_bits + + gen = torch.Generator().manual_seed(seed) + # Hyperplanes: [num_tables, num_bits, dim] — normalized for numerical + # stability (sign is unaffected, but magnitudes stay bounded). + planes = torch.randn(num_tables, num_bits, dim, generator=gen) + self.planes = F.normalize(planes, dim=-1) + + # Bit weights for packing {0,1}^B → int: 2^0, 2^1, ..., 2^(B-1). + self._bit_weights = ( + 1 << torch.arange(num_bits, dtype=torch.int64) + ) # [num_bits] + + # Populated by .build(). List of dicts (one per table): hash_int → [item_idx, ...] + self._tables: list[dict[int, list[int]]] = [{} for _ in range(num_tables)] + self._n_items = 0 + + @torch.no_grad() + def _hash(self, emb: torch.Tensor) -> torch.Tensor: + """Hash [N, D] embeddings to [N, num_tables] int64 codes. + + Project each embedding onto every hyperplane and take the sign + pattern as a bit vector; pack the bits into an int64 per table. + """ + if emb.dim() != 2: + raise ValueError(f"expected [N, D], got shape {tuple(emb.shape)}") + flat_planes = self.planes.reshape(-1, self.dim) # [T*B, D] + projections = emb @ flat_planes.t() # [N, T*B] + bits = (projections > 0).view(-1, self.num_tables, self.num_bits) # [N, T, B] + codes = (bits.to(torch.int64) * self._bit_weights).sum(dim=-1) # [N, T] + return codes + + @torch.no_grad() + def build(self, embeddings: torch.Tensor) -> None: + """Insert all [N, D] embeddings into every table. Call once per epoch.""" + if embeddings.dim() != 2 or embeddings.size(1) != self.dim: + raise ValueError( + f"embeddings must be [N, {self.dim}], got {tuple(embeddings.shape)}", + ) + for t in range(self.num_tables): + self._tables[t].clear() + codes = self._hash(embeddings.float()) # [N, num_tables] + codes_cpu = codes.cpu().tolist() + for item_idx, per_table in enumerate(codes_cpu): + for t, code in enumerate(per_table): + self._tables[t].setdefault(code, []).append(item_idx) + self._n_items = embeddings.size(0) + # Stats for logging. + bucket_sizes = [ + len(bucket) for table in self._tables for bucket in table.values() + ] + if bucket_sizes: + LOGGER.info( + "LSH built: N=%d, tables=%d, bits=%d, " + "avg bucket=%.1f (max=%d)", + self._n_items, self.num_tables, self.num_bits, + sum(bucket_sizes) / len(bucket_sizes), max(bucket_sizes), + ) + + @torch.no_grad() + def candidates(self, query: torch.Tensor) -> set[int]: + """Union of hash-bucket contents for a [D] query vector.""" + codes = self._hash(query.float().unsqueeze(0)).squeeze(0) # [num_tables] + out: set[int] = set() + codes_list = codes.tolist() + for t, code in enumerate(codes_list): + bucket = self._tables[t].get(code) + if bucket is not None: + out.update(bucket) + return out + + @property + def n_items(self) -> int: + return self._n_items diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 2c64f42..f932c05 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -39,6 +39,7 @@ from src.datasets.gtauav_dataset import ( collate_sat_gallery, ) from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler +from src.datasets.embedding_cache import EmbeddingCache from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler from src.losses.multi_infonce import InfoNCELoss from src.losses.weighted_infonce import WeightedInfoNCELoss @@ -118,6 +119,11 @@ class TrainConfigGTAUAV: sampler_type: str = "mutex" # "mutex" (no false negatives) or "dss" (DSS + mutex) dss_reembed_every: int = 1 # Re-embed train queries every N epochs for DSS. dss_warmup_epochs: int = 1 # Use mutex-only for the first N epochs (fresh model embeddings aren't useful) + dss_knn_device: str = "cuda" # Device for similarity matmul in DSS sampler. + dss_use_lsh: bool = False # Approximate kNN via LSH (opt-in; exact is fast at 25K). + dss_lsh_num_tables: int = 8 + dss_lsh_num_bits: int = 14 + dss_cache_dir: str | None = None # Disk cache for embeddings; None = disabled. # Legacy alias kept for backward compatibility. use_mutex_sampler: bool = True @@ -622,10 +628,15 @@ def train(cfg: TrainConfigGTAUAV) -> None: if effective_sampler_type == "dss": batch_sampler = DynamicSimilaritySampler( sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed, + knn_device=cfg.dss_knn_device, + use_lsh=cfg.dss_use_lsh, + lsh_num_tables=cfg.dss_lsh_num_tables, + lsh_num_bits=cfg.dss_lsh_num_bits, ) LOGGER.info( - "Sampler: DynamicSimilarity — embedding-ranked batches with mutex constraint " - "(warmup=%d epochs mutex-only, re-embed every %d epochs)", + "Sampler: DynamicSimilarity — kNN on %s%s, warmup=%d, re-embed every %d epochs", + cfg.dss_knn_device, + " + LSH" if cfg.dss_use_lsh else "", cfg.dss_warmup_epochs, cfg.dss_reembed_every, ) elif effective_sampler_type == "mutex": @@ -655,6 +666,11 @@ def train(cfg: TrainConfigGTAUAV) -> None: pin_memory=True, drop_last=True, ) + + emb_cache: EmbeddingCache | None = None + if cfg.dss_cache_dir is not None: + emb_cache = EmbeddingCache(cfg.dss_cache_dir) + LOGGER.info("DSS embedding cache: %s", cfg.dss_cache_dir) test_loader = DataLoader( test_ds, batch_size=cfg.batch_size, @@ -753,15 +769,23 @@ def train(cfg: TrainConfigGTAUAV) -> None: and epoch >= cfg.dss_warmup_epochs and (epoch - cfg.dss_warmup_epochs) % cfg.dss_reembed_every == 0 ): - LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(train_ds), epoch) - t_embed = time.time() - query_embs = _embed_drone_queries( - model, train_ds, cfg.device, - batch_size=cfg.batch_size * cfg.grad_accum_steps, - num_workers=cfg.num_workers, - ) + query_embs: torch.Tensor | None = None + if emb_cache is not None: + query_embs = emb_cache.load(epoch) + if query_embs is None: + LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(train_ds), epoch) + t_embed = time.time() + query_embs = _embed_drone_queries( + model, train_ds, cfg.device, + batch_size=cfg.batch_size * cfg.grad_accum_steps, + num_workers=cfg.num_workers, + ) + LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed) + if emb_cache is not None: + emb_cache.save(epoch, query_embs) + t_sampler = time.time() batch_sampler.update_embeddings(query_embs) - LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed) + LOGGER.info("DSS: sampler update_embeddings took %.2fs", time.time() - t_sampler) epoch_start = time.time() agg: dict[str, float] = {}