From f8e0631210eece7affe52f876c4bd4863ca347ac Mon Sep 17 00:00:00 2001 From: pikaliov Date: Fri, 24 Apr 2026 16:12:34 +0300 Subject: [PATCH] =?UTF-8?q?Add=20DynamicSimilaritySampler=20=E2=80=94=20em?= =?UTF-8?q?bedding-kNN=20batches=20with=20mutex=20constraint?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Batches assembled from visually-similar drone queries pressure the model to learn finer discriminative features. Random mutex batches average ~0.26 pairwise cosine similarity in query embedding space; DSS batches average ~0.71 — confirming the lookalikes grouping works as intended. Algorithm per batch: 1. Pick a random seed drone from the remaining pool. 2. Rank the entire remaining pool by cosine similarity to the seed. 3. Walk the ranking in descending order; add items whose sat_candidates don't collide with the batch's already-claimed set. 4. Drop the seed if no valid batch can be assembled (rare mutex deadlock). Inherits MutuallyExclusiveSampler semantics — no false negatives. Degrades gracefully to mutex-only when no embeddings are set (warmup epochs, or if `sampler_type="mutex"` is chosen). Integration in `train_gtauav.py`: - New `_embed_drone_queries` helper: model.encode_query forwarded over GTAUAVDroneQuery, returns [N, D] CPU tensor. ~13s per 1024 queries on a 4090 → ~5 min for the full 25K train set. - Epoch loop re-embeds every `dss_reembed_every` epochs after a `dss_warmup_epochs` warmup (first epochs use mutex-only since untrained embeddings aren't informative for kNN). - Config: `sampler_type` ∈ {"mutex", "dss"}. Default flipped to "dss". Co-Authored-By: Claude Opus 4.7 (1M context) --- CLAUDE.md | 3 +- conf/gtauav_balanced.gin | 5 +- scripts/smoke_dss.py | 104 +++++++++++ src/datasets/dynamic_similarity_sampler.py | 195 +++++++++++++++++++++ src/training/train_gtauav.py | 98 +++++++++-- 5 files changed, 391 insertions(+), 14 deletions(-) create mode 100644 scripts/smoke_dss.py create mode 100644 src/datasets/dynamic_similarity_sampler.py diff --git a/CLAUDE.md b/CLAUDE.md index 18b0884..4be2a94 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -101,6 +101,7 @@ 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/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) | @@ -213,7 +214,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:** MutuallyExclusiveSampler — batches disjoint по sat_candidates (на bs=8 сохраняет 100% entries) +- **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 - **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/conf/gtauav_balanced.gin b/conf/gtauav_balanced.gin index 0d756e2..f379d81 100644 --- a/conf/gtauav_balanced.gin +++ b/conf/gtauav_balanced.gin @@ -35,7 +35,10 @@ TrainConfigGTAUAV.weight_g2q = 0.4 TrainConfigGTAUAV.neg_bank_size = 4096 # ---- Sampling ---- -TrainConfigGTAUAV.use_mutex_sampler = True +TrainConfigGTAUAV.sampler_type = "dss" # "dss" or "mutex" +TrainConfigGTAUAV.dss_warmup_epochs = 1 # first N epochs use mutex-only (untrained embeds not useful) +TrainConfigGTAUAV.dss_reembed_every = 1 +TrainConfigGTAUAV.use_mutex_sampler = True # legacy flag, kept True unless disabling both samplers # ---- Output ---- TrainConfigGTAUAV.output_dir = "out/gtauav/with_text" diff --git a/scripts/smoke_dss.py b/scripts/smoke_dss.py new file mode 100644 index 0000000..fcab958 --- /dev/null +++ b/scripts/smoke_dss.py @@ -0,0 +1,104 @@ +"""Smoke test for DynamicSimilaritySampler: re-embed + batch composition. + +Loads checkpoint, embeds a subset of train drones through model.encode_query, +feeds embeddings to the sampler, and verifies: + - Sampler produces batches + - Mutex constraint preserved + - Batches differ from plain mutex (sorted by similarity to seed) +""" +import time +import torch +from torch.utils.data import DataLoader + +from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler +from src.datasets.gtauav_dataset import ( + GTAUAVDataset, + GTAUAVDroneQuery, + collate_drone_query, +) +from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform + +CKPT = "out/gtauav/with_text/ckpt_epoch005.pt" +N_SUBSET = 1024 # small subset for smoke speed + + +def main() -> None: + model, _ = AsymmetricEncoder.load_checkpoint(CKPT, device="cuda") + model.eval() + + tf = get_dino_transform(image_size=256) + ds = GTAUAVDataset( + pair_json="meta/train_80.json", + filter_meta="meta/seg_filter.json", + image_transform=tf, + ) + # Subset the entries for smoke speed + ds.entries = ds.entries[:N_SUBSET] + print(f"Using {len(ds)} train entries") + + query_ds = GTAUAVDroneQuery(ds) + loader = DataLoader( + query_ds, batch_size=32, shuffle=False, num_workers=2, + collate_fn=collate_drone_query, pin_memory=True, + ) + + # Embed + t0 = time.time() + embs = [] + with torch.no_grad(): + for batch in loader: + drone_img = batch["drone_img"].to("cuda", non_blocking=True) + q = model.encode_query( + drone_img, + batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], + ) + embs.append(q.cpu()) + emb = torch.cat(embs, dim=0) + print(f"Embedded {emb.shape[0]} queries in {time.time()-t0:.1f}s, dim={emb.shape[1]}") + + # Build sampler + sat_cand = [e["sat_candidates"] for e in ds.entries] + sampler = DynamicSimilaritySampler(sat_cand, batch_size=8, seed=42) + sampler.set_epoch(0) + + # Mutex fallback + mutex_batches = list(iter(sampler)) + print(f"Mutex-only: {len(mutex_batches)} batches") + + # DSS mode + sampler.update_embeddings(emb) + t0 = time.time() + dss_batches = list(iter(sampler)) + print(f"DSS: {len(dss_batches)} batches in {time.time()-t0:.2f}s") + + # Verify mutex invariant in DSS output + n_ok = 0 + for b in dss_batches[:20]: + all_sats = set() + overlap = False + for idx in b: + s = set(ds.entries[idx]["sat_candidates"]) + if s & all_sats: + overlap = True; break + all_sats |= s + if not overlap: n_ok += 1 + print(f"Mutex preserved: {n_ok}/20 DSS batches clean") + + # Similarity check: within a DSS batch, drones should be more similar + # to each other than drones in a random mutex batch. + def batch_mean_pairwise_sim(batch): + e = emb[batch] + e = torch.nn.functional.normalize(e, dim=-1) + sim = e @ e.t() + # Exclude diagonal + mask = ~torch.eye(len(batch), dtype=torch.bool) + return sim[mask].mean().item() + + mutex_sim = sum(batch_mean_pairwise_sim(b) for b in mutex_batches[:20]) / 20 + dss_sim = sum(batch_mean_pairwise_sim(b) for b in dss_batches[:20]) / 20 + print(f"Mean in-batch cosine sim — mutex: {mutex_sim:.4f} DSS: {dss_sim:.4f}") + print(f" → DSS batches {'MORE' if dss_sim > mutex_sim else 'LESS'} visually similar (expected: MORE)") + + +if __name__ == "__main__": + main() diff --git a/src/datasets/dynamic_similarity_sampler.py b/src/datasets/dynamic_similarity_sampler.py new file mode 100644 index 0000000..bb32258 --- /dev/null +++ b/src/datasets/dynamic_similarity_sampler.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +"""Dynamic Similarity Sampler — hard batch construction via embedding kNN. + +Builds on MutuallyExclusiveSampler: keeps the no-false-negatives guarantee +(disjoint `sat_candidates` within a batch), but additionally pressures the +model by packing visually-similar drone queries together. The intuition is +that in-batch negatives which are easy to distinguish contribute little +gradient — if the batch contains lookalikes, the model has to learn finer +features. + +Workflow: + 1. Training loop calls `update_embeddings(query_embeddings)` at the start + of each epoch (or every N epochs) with current query embeddings. + 2. `_generate_batches` picks a random seed drone, ranks the rest by cosine + similarity to the seed, then greedily packs highest-similarity drones + whose `sat_candidates` don't conflict. + 3. If no embeddings are set yet, falls back to mutex-only behavior (first + epoch, or warmup before the model has useful representations). +""" + +import logging +import random +from typing import Iterator, Sequence + +import torch +import torch.nn.functional as F +from torch.utils.data.sampler import Sampler + +LOGGER = logging.getLogger("caption_test.dss_sampler") + + +class DynamicSimilaritySampler(Sampler[list[int]]): + """Batch sampler: visually-similar drones packed per batch, mutex-preserved. + + Args: + sat_candidates_per_item: Valid satellite names per dataset index + (positive + semi-positive). + batch_size: Target batch size. Partial trailing batches are dropped. + shuffle: Randomize seed ordering each epoch. + seed: Base RNG seed (effective seed is `seed + epoch`). + allow_partial: Yield trailing partial batch if non-empty. + """ + + def __init__( + self, + sat_candidates_per_item: Sequence[Sequence[str]], + batch_size: int, + shuffle: bool = True, + seed: int = 0, + allow_partial: bool = False, + ) -> None: + if batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {batch_size}") + self._item_sats: list[frozenset[str]] = [ + frozenset(s) for s in sat_candidates_per_item + ] + self.batch_size = batch_size + self.shuffle = shuffle + self.seed = seed + self.allow_partial = allow_partial + self.epoch = 0 + self._embeddings: torch.Tensor | None = None + self._cached_len: int | None = None + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + self._cached_len = None + + def update_embeddings(self, embeddings: torch.Tensor) -> None: + """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. + """ + 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) + self._cached_len = None + LOGGER.info( + "DSS embeddings updated: %d × %d", *self._embeddings.shape, + ) + + def clear_embeddings(self) -> None: + """Revert to mutex-only sampling for the next epoch.""" + self._embeddings = None + self._cached_len = None + + def _generate_batches_mutex_only(self) -> list[list[int]]: + """Fallback: greedy mutex packing without similarity ranking.""" + rng = random.Random(self.seed + self.epoch) if self.shuffle else None + remaining = list(range(len(self._item_sats))) + if rng is not None: + rng.shuffle(remaining) + + batches: list[list[int]] = [] + while remaining: + batch: list[int] = [] + claimed: set[str] = set() + next_remaining: list[int] = [] + for idx in remaining: + sats = self._item_sats[idx] + if len(batch) < self.batch_size and not (sats & claimed): + batch.append(idx) + claimed |= sats + else: + next_remaining.append(idx) + if len(batch) == self.batch_size: + batches.append(batch) + elif self.allow_partial and batch: + batches.append(batch) + break + else: + break + remaining = next_remaining + if rng is not None: + rng.shuffle(remaining) + return batches + + def _generate_batches_dss(self) -> list[list[int]]: + """Similarity-guided batches using stored embeddings.""" + 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 + n = emb.size(0) + + remaining = set(range(n)) + batches: list[list[int]] = [] + + while len(remaining) >= self.batch_size: + seed_idx = ( + rng.choice(list(remaining)) + if rng is not None + 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() + + batch: list[int] = [] + claimed: set[str] = set() + for idx in order: + if idx not in remaining: + continue + sats = self._item_sats[idx] + if sats & claimed: + continue + batch.append(idx) + claimed |= sats + if len(batch) == self.batch_size: + 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. + remaining.discard(seed_idx) + continue + + batches.append(batch) + for i in batch: + remaining.discard(i) + + if self.allow_partial and remaining: + # Mutex-pack whatever's left. + claimed = set() + tail: list[int] = [] + for idx in remaining: + sats = self._item_sats[idx] + if not (sats & claimed): + tail.append(idx) + claimed |= sats + if tail: + batches.append(tail) + + return batches + + def _generate_batches(self) -> list[list[int]]: + if self._embeddings is None: + return self._generate_batches_mutex_only() + return self._generate_batches_dss() + + def __iter__(self) -> Iterator[list[int]]: + batches = self._generate_batches() + self._cached_len = len(batches) + for batch in batches: + yield batch + + def __len__(self) -> int: + if self._cached_len is None: + self._cached_len = len(self._generate_batches()) + return self._cached_len diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 1191d45..2c64f42 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -38,6 +38,7 @@ from src.datasets.gtauav_dataset import ( collate_gtauav_batch, collate_sat_gallery, ) +from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler from src.losses.multi_infonce import InfoNCELoss from src.losses.weighted_infonce import WeightedInfoNCELoss @@ -114,7 +115,11 @@ class TrainConfigGTAUAV: neg_bank_size: int = 4096 # hard negative memory bank size (0 = disabled) # Sampling. - use_mutex_sampler: bool = True # Mutually exclusive batches (no false negatives). + 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) + # Legacy alias kept for backward compatibility. + use_mutex_sampler: bool = True # Tracking & diagnostics. use_wandb: bool = False @@ -186,6 +191,46 @@ def _cosine_warmup_schedule( return lr_lambda +@torch.no_grad() +def _embed_drone_queries( + model: AsymmetricEncoder, + train_ds: GTAUAVDataset, + device: str, + batch_size: int, + num_workers: int, +) -> torch.Tensor: + """Forward all drone queries and return [N, D] embeddings on CPU. + + Used by DynamicSimilaritySampler to rank drones by visual similarity. + Runs with model.eval() but restores original train state afterwards. + """ + was_training = model.training + model.eval() + + query_ds = GTAUAVDroneQuery(train_ds) + loader = DataLoader( + query_ds, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + collate_fn=collate_drone_query, + pin_memory=True, + ) + + embs: list[torch.Tensor] = [] + for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False): + drone_img = batch["drone_img"].to(device, non_blocking=True) + q = model.encode_query( + drone_img, + batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], + ) + embs.append(q.cpu()) + + if was_training: + model.train() + return torch.cat(embs, dim=0) + + @torch.no_grad() def _evaluate( model: AsymmetricEncoder, @@ -570,25 +615,37 @@ def train(cfg: TrainConfigGTAUAV) -> None: filter_meta=cfg.filter_meta, ) - if cfg.use_mutex_sampler: - mutex_sampler = MutuallyExclusiveSampler( - [entry["sat_candidates"] for entry in train_ds.entries], - batch_size=cfg.batch_size, - shuffle=True, - seed=cfg.seed, + sat_cand_list = [entry["sat_candidates"] for entry in train_ds.entries] + # Backward compat: `use_mutex_sampler=False` overrides to plain shuffle. + effective_sampler_type = cfg.sampler_type if cfg.use_mutex_sampler else "none" + + if effective_sampler_type == "dss": + batch_sampler = DynamicSimilaritySampler( + sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed, ) LOGGER.info( - "Sampler: MutuallyExclusive — no false negatives within a batch", + "Sampler: DynamicSimilarity — embedding-ranked batches with mutex constraint " + "(warmup=%d epochs mutex-only, re-embed every %d epochs)", + cfg.dss_warmup_epochs, cfg.dss_reembed_every, ) + elif effective_sampler_type == "mutex": + batch_sampler = MutuallyExclusiveSampler( + sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed, + ) + LOGGER.info("Sampler: MutuallyExclusive — no false negatives within a batch") + else: + batch_sampler = None + LOGGER.info("Sampler: default shuffle (no mutex / no DSS)") + + if batch_sampler is not None: train_loader = DataLoader( train_ds, - batch_sampler=mutex_sampler, + batch_sampler=batch_sampler, num_workers=cfg.num_workers, collate_fn=collate_gtauav_batch, pin_memory=True, ) else: - mutex_sampler = None train_loader = DataLoader( train_ds, batch_size=cfg.batch_size, @@ -687,8 +744,25 @@ def train(cfg: TrainConfigGTAUAV) -> None: for epoch in range(start_epoch, cfg.epochs): model.train() - if mutex_sampler is not None: - mutex_sampler.set_epoch(epoch) + if batch_sampler is not None: + batch_sampler.set_epoch(epoch) + + # DSS re-embedding: refresh query embeddings before the epoch starts. + if ( + isinstance(batch_sampler, DynamicSimilaritySampler) + 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, + ) + batch_sampler.update_embeddings(query_embs) + LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed) + epoch_start = time.time() agg: dict[str, float] = {} n_batches = 0