f8e0631210
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) <noreply@anthropic.com>
196 lines
7.0 KiB
Python
196 lines
7.0 KiB
Python
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
|