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>
105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
"""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()
|