diff --git a/CLAUDE.md b/CLAUDE.md index 49875a3..5028961 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -30,9 +30,13 @@ TextFusionMLP shared между query и gallery (одинаковый форм Для sat images без captions: s_txt=None → g = s_img (gate passthrough) LOSS: L = 0.6·CE(q̂·ĝᵀ/τ, targets) + 0.4·CE(ĝ·q̂ᵀ/τ, targets) - τ = 1/exp(logit_scale), learnable, clamped [0.01, 0.5], init=0.07 + τ = 1/exp(logit_scale), learnable, clamped [0.01, 0.1], init=0.07 label_smoothing=0.1 +BATCH SAMPLING: MutuallyExclusiveSampler — в одном батче нет двух drone'ов + с пересекающимися sat_candidates (исключает false negatives, которые + иначе появляются из-за multi-positive структуры GTA-UAV). + BASELINE: σ(α_q)=σ(α_g)=1.0, text disabled, DGTRS not loaded ``` @@ -42,7 +46,7 @@ BASELINE: σ(α_q)=σ(α_g)=1.0, text disabled, DGTRS not loaded - **L3 fingerprint:** P3 — уникальные landmarks для matching (20-50 tok) - **Fusion:** z_text = MLP([z₁; z₂; z₃]) — concat 3×768 → Linear(2304,1024) → GELU → Linear(1024,1024) - **Shared MLP** между query и gallery ветками (одинаковый формат captions) -- **Satellite captions:** 6,546 из 14,640 sat images имеют captions. Для остальных gate passthrough (g = s_img) +- **Satellite captions:** 6,546 из 14,640 sat images имеют captions. Для остальных gate passthrough (g = s_img) — **per-sample mask** в `_fuse_with_mask` возвращает чистые image features для samples без caption (без шума от пустых строк) ### Text encoder: DGTRS-CLIP (official architecture) - Код: `src/models/dgtrs/` — из github.com/MitsuiChen14/DGTRS (Apache-2.0) @@ -94,10 +98,14 @@ Eval: Resize(256) + CenterCrop(256) + ImageNet normalization. |------|-----------| | `src/models/dgtrs/model.py` | Официальная архитектура DGTRS-CLIP text encoder (Apache-2.0) | | `src/models/dgtrs/simple_tokenizer.py` | BPE tokenizer (248 tokens, vocab 49408) | -| `src/models/asymmetric_encoder.py` | DINOv3ViT + TextFusionMLP + AsymmetricEncoder + GatedFusion | -| `src/datasets/gtauav_dataset.py` | GTA-UAV-LR loader + L1/L2/L3 caption parsing из VLM JSON | -| `src/losses/multi_infonce.py` | InfoNCE с learnable temperature (fp32), clamp [0.01, 0.5] | -| `src/training/train_gtauav.py` | Training loop с gin, W&B/TB, AMP, per-group LR, warmup, --resume | +| `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/losses/multi_infonce.py` | **Primary:** SymmetricInfoNCE + MoCo queue, learnable τ clamp [0.01, 0.1], weights q2g=0.6 g2q=0.4 | +| `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) | +| `src/training/train_gtauav.py` | Training loop: full-gallery `_evaluate`, mutex sampler wiring, loss_type switch | +| `scripts/smoke_eval.py` / `scripts/smoke_train.py` | Регрессионные smoke-тесты для eval и train pipeline | | `src/training/trackers.py` | Unified experiment tracker: W&B + TensorBoard + CSV | | `src/training/grad_monitor.py` | Gradient norm monitoring per param group | | `src/training/gradcam.py` | Grad-CAM visualization для DINOv3 encoders | @@ -203,7 +211,9 @@ Meta-файл `meta/seg_filter.json`: исключение изображени - 10 epochs, batch 64, AMP, image 256x256 - **Optimizer:** AdamW, per-group LR: proj=1e-4, text=1e-5 (10x lower) - **Scheduler:** linear warmup (2 epochs) + cosine annealing (per-step) -- **Loss:** InfoNCE с learnable temperature (CLIP logit_scale), init=0.07, clamp [0.01, 0.5] +- **Loss:** SymmetricInfoNCE (q2g=0.6, g2q=0.4) с learnable τ (init=0.07, clamp [0.01, 0.1]) +- **Batch sampler:** MutuallyExclusiveSampler — batches disjoint по sat_candidates (на bs=8 сохраняет 100% entries) +- **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 - Satellite: RandomResizedCrop(0.7-1.0), HFlip, ColorJitter, Grayscale(5%) diff --git a/conf/gtauav_balanced.gin b/conf/gtauav_balanced.gin index 7a5247d..9461348 100644 --- a/conf/gtauav_balanced.gin +++ b/conf/gtauav_balanced.gin @@ -1,8 +1,8 @@ # GTA-UAV Balanced: Asymmetric DINOv3 (WEB+SAT) with L1/L2/L3 captions. -# WeightedInfoNCE loss for GTA-UAV partial overlap handling. +# Symmetric InfoNCE + MutuallyExclusiveSampler (no false negatives). # 10 epochs, MONA all 24 blocks, 1024-dim retrieval, hard negative bank. -import src.losses.weighted_infonce +import src.losses.multi_infonce # ---- Training ---- TrainConfigGTAUAV.epochs = 10 @@ -26,11 +26,17 @@ TrainConfigGTAUAV.shared_encoder = False TrainConfigGTAUAV.gradient_checkpointing = True # ---- Loss ---- +TrainConfigGTAUAV.loss_type = "symmetric" TrainConfigGTAUAV.tau_init = 0.07 TrainConfigGTAUAV.label_smoothing = 0.1 TrainConfigGTAUAV.learnable_temperature = True +TrainConfigGTAUAV.weight_q2g = 0.6 +TrainConfigGTAUAV.weight_g2q = 0.4 TrainConfigGTAUAV.neg_bank_size = 4096 +# ---- Sampling ---- +TrainConfigGTAUAV.use_mutex_sampler = True + # ---- Output ---- TrainConfigGTAUAV.output_dir = "out/gtauav/with_text" @@ -42,8 +48,11 @@ TrainConfigGTAUAV.gradcam_every = 5 TrainConfigGTAUAV.use_profiler = False TrainConfigGTAUAV.log_grad_norms = True -# ---- WeightedInfoNCE (gin-configurable) ---- -WeightedInfoNCELoss.temperature_init = 0.07 -WeightedInfoNCELoss.learnable_temperature = True -WeightedInfoNCELoss.label_smoothing = 0.1 -WeightedInfoNCELoss.k = 5.0 +# ---- InfoNCELoss (gin-configurable) ---- +InfoNCELoss.temperature_init = 0.07 +InfoNCELoss.learnable_temperature = True +InfoNCELoss.label_smoothing = 0.1 +InfoNCELoss.weight_q2g = 0.6 +InfoNCELoss.weight_g2q = 0.4 +InfoNCELoss.tau_min = 0.01 +InfoNCELoss.tau_max = 0.1 diff --git a/scripts/smoke_eval.py b/scripts/smoke_eval.py new file mode 100644 index 0000000..449abc2 --- /dev/null +++ b/scripts/smoke_eval.py @@ -0,0 +1,46 @@ +"""Smoke-test for the rewritten `_evaluate` function. + +Loads checkpoint ckpt_epoch005.pt and runs the new full-gallery eval with +max_batches=5 to verify end-to-end without waiting for a full epoch. +""" +from torch.utils.data import DataLoader + +from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch +from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform +from src.training.train_gtauav import _evaluate + +CKPT = "out/gtauav/with_text/ckpt_epoch005.pt" + +def main() -> None: + model, _ = AsymmetricEncoder.load_checkpoint(CKPT, device="cuda") + + eval_tf = get_dino_transform(image_size=256) + ds = GTAUAVDataset( + pair_json="meta/test_20.json", + filter_meta="meta/seg_filter.json", + image_transform=eval_tf, + ) + loader = DataLoader( + ds, + batch_size=32, + shuffle=False, + num_workers=2, + collate_fn=collate_gtauav_batch, + pin_memory=True, + ) + + print("Running _evaluate (max_batches=5 on queries, full gallery)...") + metrics = _evaluate( + model=model, + loader=loader, + device="cuda", + max_batches=5, + desc="smoke", + ) + print("--- metrics ---") + for k, v in metrics.items(): + print(f" {k}: {v:.4f}" if isinstance(v, float) else f" {k}: {v}") + + +if __name__ == "__main__": + main() diff --git a/scripts/smoke_train.py b/scripts/smoke_train.py new file mode 100644 index 0000000..3270943 --- /dev/null +++ b/scripts/smoke_train.py @@ -0,0 +1,79 @@ +"""Minimal training smoke test: 2 batches forward+backward. + +Verifies end-to-end that MutuallyExclusiveSampler + InfoNCELoss + +per-sample caption masking compose correctly for training. +""" +import torch +from torch.utils.data import DataLoader + +from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch +from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler +from src.losses.multi_infonce import InfoNCELoss +from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform + +CKPT = "out/gtauav/with_text/ckpt_epoch005.pt" + + +def main() -> None: + model, _ = AsymmetricEncoder.load_checkpoint(CKPT, device="cuda") + model.train() + + tf = get_dino_transform(image_size=256) + ds = GTAUAVDataset( + pair_json="meta/train_80.json", + filter_meta="meta/seg_filter.json", + drone_transform=tf, + sat_transform=tf, + ) + + sampler = MutuallyExclusiveSampler( + [e["sat_candidates"] for e in ds.entries], + batch_size=8, shuffle=True, seed=42, + ) + sampler.set_epoch(0) + loader = DataLoader( + ds, batch_sampler=sampler, num_workers=2, + collate_fn=collate_gtauav_batch, pin_memory=True, + ) + + loss_fn = InfoNCELoss( + temperature_init=0.07, learnable_temperature=True, + label_smoothing=0.1, weight_q2g=0.6, weight_g2q=0.4, + tau_min=0.01, tau_max=0.1, + ).to("cuda") + + trainable = [p for p in model.trainable_parameters()] + list(loss_fn.parameters()) + opt = torch.optim.AdamW(trainable, lr=1e-4) + + it = iter(loader) + for step in range(2): + batch = next(it) + opt.zero_grad() + emb = model( + drone_img=batch["drone_img"].to("cuda", non_blocking=True), + sat_img=batch["sat_img"].to("cuda", non_blocking=True), + caption_l1=batch["caption_l1"], caption_l2=batch["caption_l2"], caption_l3=batch["caption_l3"], + sat_caption_l1=batch["sat_caption_l1"], sat_caption_l2=batch["sat_caption_l2"], sat_caption_l3=batch["sat_caption_l3"], + ) + out = loss_fn(emb, epoch=0, total_epochs=10) + out["total"].backward() + opt.step() + + # Verify mutual exclusion in batch + batch_sats = [set(ds.entries[i]["sat_candidates"]) for i in batch.get("__indices__", range(8))] + # We can also check via sat_names (one sat per drone sampled) + sat_names = batch["sat_names"] + print( + f" step {step}: loss={out['total'].item():.4f} " + f"tau={out['temperature'].item():.4f} " + f"gate_q={out['gate_q'].item():.3f} gate_g={out['gate_g'].item():.3f} " + f"n_drone_caps={sum(1 for t in batch['caption_l1'] if t)} " + f"n_sat_caps={sum(1 for t in batch['sat_caption_l1'] if t)}" + ) + assert torch.isfinite(out["total"]).all(), "Loss not finite!" + + print("OK: 2 train steps completed with finite loss") + + +if __name__ == "__main__": + main() diff --git a/src/datasets/gtauav_dataset.py b/src/datasets/gtauav_dataset.py index 3a0f02a..10b0c75 100644 --- a/src/datasets/gtauav_dataset.py +++ b/src/datasets/gtauav_dataset.py @@ -287,3 +287,109 @@ def collate_gtauav_batch( "sat_names": [b["sat_name"] for b in batch], "positive_weights": torch.tensor([b["positive_weight"] for b in batch], dtype=torch.float32), } + + +def _load_rgb_image( + rgb_root: Path, + directory: str, + filename: str, + transform: Callable | None, +) -> torch.Tensor: + path = rgb_root / directory / filename + with Image.open(path) as img: + rgb = img.convert("RGB") + if transform is not None: + return transform(rgb) + return torch.tensor(0) + + +class GTAUAVSatGallery(Dataset): + """Unique satellite gallery for retrieval evaluation. + + Takes a GTAUAVDataset and extracts the set of unique satellite names + appearing in any entry's sat_candidates. Yields one (sat_img, captions) + per unique name — suitable as the gallery side of a retrieval benchmark. + + Used by `_evaluate` in train_gtauav.py to forward the full gallery once. + """ + + def __init__(self, source: "GTAUAVDataset") -> None: + self.rgb_root = source.rgb_root + self.sat_transform = source.sat_transform + # Collect unique sats (preserve first-seen order for determinism). + unique: dict[str, tuple[str, tuple[str, str, str] | None]] = {} + for entry in source.entries: + sat_dir = entry["sat_dir"] + for sat_name in entry["sat_candidates"]: + if sat_name in unique: + continue + caps = entry["sat_captions"].get(sat_name) + unique[sat_name] = (sat_dir, caps) + self.sat_names: list[str] = list(unique.keys()) + self._sat_info: dict[str, tuple[str, tuple[str, str, str] | None]] = unique + + def __len__(self) -> int: + return len(self.sat_names) + + def __getitem__(self, idx: int) -> dict[str, Any]: + sat_name = self.sat_names[idx] + sat_dir, caps = self._sat_info[sat_name] + sat_img = _load_rgb_image(self.rgb_root, sat_dir, sat_name, self.sat_transform) + if caps is not None: + l1, l2, l3 = caps + else: + l1 = l2 = l3 = _EMPTY_CAPTION + return { + "sat_img": sat_img, + "sat_name": sat_name, + "sat_caption_l1": l1, + "sat_caption_l2": l2, + "sat_caption_l3": l3, + } + + +class GTAUAVDroneQuery(Dataset): + """Drone queries with valid satellite names for multi-match evaluation.""" + + def __init__(self, source: "GTAUAVDataset") -> None: + self.rgb_root = source.rgb_root + self.drone_transform = source.drone_transform + self.entries = source.entries + + def __len__(self) -> int: + return len(self.entries) + + def __getitem__(self, idx: int) -> dict[str, Any]: + entry = self.entries[idx] + drone_img = _load_rgb_image( + self.rgb_root, entry["drone_dir"], entry["drone_name"], self.drone_transform, + ) + return { + "drone_img": drone_img, + "drone_name": entry["drone_name"], + "caption_l1": entry["caption_l1"], + "caption_l2": entry["caption_l2"], + "caption_l3": entry["caption_l3"], + "valid_sat_names": list(entry["sat_candidates"]), + } + + +def collate_sat_gallery(batch: list[dict[str, Any]]) -> dict[str, Any]: + return { + "sat_img": torch.stack([b["sat_img"] for b in batch], dim=0), + "sat_names": [b["sat_name"] for b in batch], + "sat_caption_l1": [b["sat_caption_l1"] for b in batch], + "sat_caption_l2": [b["sat_caption_l2"] for b in batch], + "sat_caption_l3": [b["sat_caption_l3"] for b in batch], + } + + +def collate_drone_query(batch: list[dict[str, Any]]) -> dict[str, Any]: + return { + "drone_img": torch.stack([b["drone_img"] for b in batch], dim=0), + "drone_names": [b["drone_name"] for b in batch], + "caption_l1": [b["caption_l1"] for b in batch], + "caption_l2": [b["caption_l2"] for b in batch], + "caption_l3": [b["caption_l3"] for b in batch], + "valid_sat_names": [b["valid_sat_names"] for b in batch], + } diff --git a/src/datasets/mutually_exclusive_sampler.py b/src/datasets/mutually_exclusive_sampler.py new file mode 100644 index 0000000..ac03cb9 --- /dev/null +++ b/src/datasets/mutually_exclusive_sampler.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +"""Batch sampler that prevents false negatives from GTA-UAV's semi-positive graph. + +In GTA-UAV a single satellite tile can be a valid (semi-)positive for multiple +drone frames. When those frames land in the same InfoNCE batch, the shared +satellite becomes a negative for every drone except one — which trains the +model to push apart embeddings that should actually be close. + +MutuallyExclusiveSampler resolves this by greedily building batches where no +two drone indices share ANY entry in their `sat_candidates` set. This keeps +the diagonal InfoNCE formulation valid: every off-diagonal satellite is a +genuine negative for every query in the row/column. + +References: Zhu et al., Game4Loc (arXiv:2409.16925), §3.2 "Sample ID". +""" + +import logging +import random +from typing import Iterator, Sequence + +from torch.utils.data.sampler import Sampler + +LOGGER = logging.getLogger("caption_test.mutex_sampler") + + +class MutuallyExclusiveSampler(Sampler[list[int]]): + """Batch sampler yielding drone-index lists with disjoint sat_candidates. + + Args: + sat_candidates_per_item: For each dataset index i, the list of + satellite names considered valid matches (positive + semi-positive). + batch_size: Target batch size. Partial batches are dropped. + shuffle: Shuffle the index pool each epoch (use set_epoch for reproducibility). + seed: Base RNG seed — the effective seed is `seed + epoch`. + allow_partial: If True, yield the trailing partial batch. Default False. + """ + + 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._cached_len: int | None = None + + def set_epoch(self, epoch: int) -> None: + """Advance epoch (invalidates cached length) — call from training loop.""" + self.epoch = epoch + self._cached_len = None + + def _generate_batches(self) -> list[list[int]]: + 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]] = [] + + # Each outer iteration produces at most one batch. Items that conflict + # with the batch's claimed-sat set roll over to the next iteration. + 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 # leftover rollover produces no more full batches + else: + break # drop partial trailing batch + + remaining = next_remaining + if rng is not None: + rng.shuffle(remaining) + + return batches + + 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: + # Estimate by actually generating — correct count needed by + # DataLoader/tqdm. Cached until next set_epoch. + self._cached_len = len(self._generate_batches()) + return self._cached_len diff --git a/src/losses/multi_infonce.py b/src/losses/multi_infonce.py index 1ff4d06..87a3449 100644 --- a/src/losses/multi_infonce.py +++ b/src/losses/multi_infonce.py @@ -94,7 +94,7 @@ class InfoNCELoss(nn.Module): weight_g2q: float = 0.4, learnable_temperature: bool = True, tau_min: float = 0.01, - tau_max: float = 0.5, + tau_max: float = 0.1, ) -> None: super().__init__() self.temperature_init = temperature_init @@ -144,7 +144,7 @@ class InfoNCELoss(nn.Module): if self.learnable_temperature: # Clamp logit_scale in logit space first to prevent exp() overflow in fp16. # tau_min=0.01 -> max logit_scale=ln(1/0.01)=4.6 - # tau_max=0.5 -> min logit_scale=ln(1/0.5)=0.69 + # tau_max=0.1 -> min logit_scale=ln(1/0.1)=2.30 clamped = self.logit_scale.float().clamp( min=math.log(1.0 / self.tau_max), max=math.log(1.0 / self.tau_min), diff --git a/src/losses/weighted_infonce.py b/src/losses/weighted_infonce.py index 7148a04..743d10d 100644 --- a/src/losses/weighted_infonce.py +++ b/src/losses/weighted_infonce.py @@ -45,7 +45,7 @@ class WeightedInfoNCELoss(nn.Module): label_smoothing: float = 0.1, k: float = 5.0, tau_min: float = 0.01, - tau_max: float = 0.5, + tau_max: float = 0.1, ) -> None: super().__init__() self.label_smoothing = label_smoothing diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 9f077a6..23b63cd 100644 --- a/src/models/asymmetric_encoder.py +++ b/src/models/asymmetric_encoder.py @@ -371,7 +371,8 @@ class AsymmetricEncoder(nn.Module): Returns None if all captions are empty (no text available). For mixed batches (some have captions, some don't), encodes all - and lets GatedFusion handle per-sample gating. + texts (empty strings tokenize to pad+EOS — their outputs must be + masked downstream, see `_fuse_with_mask`). """ # Check if any caption is non-empty. if all(t == "" for t in l1_texts): @@ -388,6 +389,74 @@ class AsymmetricEncoder(nn.Module): tokens = tokenize_dgtrs(list(texts)).to(self.device) return self.text_encoder(tokens) + def _fuse_with_mask( + self, + img_feat: torch.Tensor, + l1_texts: list[str] | None, + l2_texts: list[str] | None, + l3_texts: list[str] | None, + fusion: GatedFusion, + ) -> torch.Tensor: + """Fuse image features with optional text, respecting per-sample presence. + + For samples where caption is an empty string, output falls back to + pure image features (avoiding noise contamination from empty-string + text embeddings). For samples with captions, applies the standard + gated fusion `σ(α)·img + (1-σ(α))·text`. + + Returns L2-normalized [B, D] embedding. + """ + if ( + self.baseline_mode + or l1_texts is None + or l2_texts is None + or l3_texts is None + ): + return F.normalize(fusion(img_feat, None), dim=-1) + + has_text = torch.tensor( + [t != "" for t in l1_texts], dtype=torch.bool, device=img_feat.device, + ) + if not has_text.any(): + return F.normalize(fusion(img_feat, None), dim=-1) + + z_text = self.encode_text_levels(l1_texts, l2_texts, l3_texts) + if z_text is None: + return F.normalize(fusion(img_feat, None), dim=-1) + + # Per-sample fusion: text-present samples use full gated fusion, + # empty-caption samples pass through pure image features. + gate = torch.sigmoid(fusion.alpha) + fused_with_text = gate * img_feat + (1.0 - gate) * z_text + out = torch.where(has_text.unsqueeze(-1), fused_with_text, img_feat) + return F.normalize(out, dim=-1) + + def encode_query( + self, + drone_img: torch.Tensor, + caption_l1: list[str] | None = None, + caption_l2: list[str] | None = None, + caption_l3: list[str] | None = None, + ) -> torch.Tensor: + """Encode drone → normalized query embedding with per-sample text mask.""" + drone_feat = self.encode_drone(drone_img) + return self._fuse_with_mask( + drone_feat, caption_l1, caption_l2, caption_l3, self.fusion_query, + ) + + def encode_gallery( + self, + sat_img: torch.Tensor, + sat_caption_l1: list[str] | None = None, + sat_caption_l2: list[str] | None = None, + sat_caption_l3: list[str] | None = None, + ) -> torch.Tensor: + """Encode satellite → normalized gallery embedding with per-sample text mask.""" + sat_feat = self.encode_satellite(sat_img) + return self._fuse_with_mask( + sat_feat, sat_caption_l1, sat_caption_l2, sat_caption_l3, self.fusion_gallery, + ) + def forward( self, drone_img: torch.Tensor, @@ -401,6 +470,10 @@ class AsymmetricEncoder(nn.Module): ) -> dict[str, torch.Tensor]: """Forward pass. + Both branches use per-sample caption masking: samples with an empty + caption string fall back to pure image features instead of being + fused with noise from empty-string text embeddings. + Args: drone_img: Drone images [B, 3, 256, 256]. sat_img: Satellite images [B, 3, 256, 256]. @@ -411,28 +484,8 @@ class AsymmetricEncoder(nn.Module): Dict with 'query' [B, embed_dim], 'gallery' [B, embed_dim], 'gate_q', 'gate_g'. """ - # Image features (frozen DINOv3). - drone_feat = self.encode_drone(drone_img) - sat_feat = self.encode_satellite(sat_img) - - # Query branch: drone + drone text. - drone_text = None - if (caption_l1 is not None and caption_l2 is not None - and caption_l3 is not None and not self.baseline_mode): - drone_text = self.encode_text_levels(caption_l1, caption_l2, caption_l3) - - query = self.fusion_query(drone_feat, drone_text) - query = F.normalize(query, dim=-1) - - # Gallery branch: satellite + satellite text. - sat_text = None - if (sat_caption_l1 is not None and sat_caption_l2 is not None - and sat_caption_l3 is not None and not self.baseline_mode): - sat_text = self.encode_text_levels(sat_caption_l1, sat_caption_l2, sat_caption_l3) - - gallery = self.fusion_gallery(sat_feat, sat_text) - gallery = F.normalize(gallery, dim=-1) - + query = self.encode_query(drone_img, caption_l1, caption_l2, caption_l3) + gallery = self.encode_gallery(sat_img, sat_caption_l1, sat_caption_l2, sat_caption_l3) return { "query": query, "gallery": gallery, diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 492dc91..1191d45 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -23,13 +23,23 @@ import gin import pandas as pd import torch import torch.nn as nn +import torch.nn.functional as F from torch.amp import GradScaler, autocast from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR from torch.utils.data import DataLoader from tqdm import tqdm -from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch +from src.datasets.gtauav_dataset import ( + GTAUAVDataset, + GTAUAVDroneQuery, + GTAUAVSatGallery, + collate_drone_query, + collate_gtauav_batch, + collate_sat_gallery, +) +from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler +from src.losses.multi_infonce import InfoNCELoss from src.losses.weighted_infonce import WeightedInfoNCELoss from src.losses.hard_negatives import NegativeMemoryBank from src.training.plot_metrics import generate_plots @@ -95,11 +105,17 @@ class TrainConfigGTAUAV: device: str = "cuda" # Loss. + loss_type: str = "symmetric" # "symmetric" (InfoNCE) or "weighted" (WeightedInfoNCE) tau_init: float = 0.07 label_smoothing: float = 0.1 learnable_temperature: bool = True + weight_q2g: float = 0.6 + weight_g2q: float = 0.4 neg_bank_size: int = 4096 # hard negative memory bank size (0 = disabled) + # Sampling. + use_mutex_sampler: bool = True # Mutually exclusive batches (no false negatives). + # Tracking & diagnostics. use_wandb: bool = False use_tb: bool = True @@ -182,109 +198,139 @@ def _evaluate( max_batches: int | None = None, desc: str = "eval", ) -> dict[str, float]: - """Compute R@K with multi-match support for GTA-UAV. + """Compute R@K and MRR on the full satellite gallery. - GTA-UAV has partial overlap between satellite crops — multiple satellites - can be valid matches for one drone. We build a valid_matches list from - the dataset and check if ANY valid match is in top-K (not just diagonal). + Standard CVGL retrieval: forward every unique satellite in the dataset + once (gallery), forward every drone query, then rank gallery by + cosine similarity. A query counts as a hit@K if ANY of its valid + satellite matches (pair_pos_sate_img_list ∪ pair_pos_semipos_sate_img_list) + appears in the top-K. + + `max_batches` subsamples the drone queries (not the gallery) — useful + for a quick train-side sanity check. """ - model.eval() - all_query: list[torch.Tensor] = [] - all_gallery: list[torch.Tensor] = [] - all_sat_names: list[str] = [] - batch_losses: list[float] = [] + dataset = loader.dataset + if not isinstance(dataset, GTAUAVDataset): + raise TypeError(f"_evaluate expects GTAUAVDataset, got {type(dataset).__name__}") - for i, batch in enumerate(tqdm(loader, desc=f" {desc}", unit="batch", leave=False)): + model.eval() + + batch_size = loader.batch_size or 32 + num_workers = getattr(loader, "num_workers", 0) + pin_memory = getattr(loader, "pin_memory", False) + + gallery_ds = GTAUAVSatGallery(dataset) + query_ds = GTAUAVDroneQuery(dataset) + + gallery_loader = DataLoader( + gallery_ds, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + pin_memory=pin_memory, + collate_fn=collate_sat_gallery, + ) + query_loader = DataLoader( + query_ds, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + pin_memory=pin_memory, + collate_fn=collate_drone_query, + ) + + # --- Gallery forward (all unique sats) --- + gallery_embs: list[torch.Tensor] = [] + gallery_names: list[str] = [] + for batch in tqdm(gallery_loader, desc=f" {desc}-gallery", unit="batch", leave=False): + sat_img = batch["sat_img"].to(device, non_blocking=True) + g = model.encode_gallery( + sat_img, + batch["sat_caption_l1"], batch["sat_caption_l2"], batch["sat_caption_l3"], + ) + gallery_embs.append(g.cpu()) + gallery_names.extend(batch["sat_names"]) + gallery = torch.cat(gallery_embs, dim=0) # [N_sat, D] + + # --- Query forward (optionally subsampled via max_batches) --- + query_embs: list[torch.Tensor] = [] + query_valid_names: list[list[str]] = [] + batch_losses: list[float] = [] + sat_name_to_idx: dict[str, int] = {name: i for i, name in enumerate(gallery_names)} + + for i, batch in enumerate(tqdm(query_loader, desc=f" {desc}-query", unit="batch", leave=False)): if max_batches is not None and i >= max_batches: break drone_img = batch["drone_img"].to(device, non_blocking=True) - sat_img = batch["sat_img"].to(device, non_blocking=True) + q = model.encode_query( + drone_img, + batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], + ) + query_embs.append(q.cpu()) + query_valid_names.extend(batch["valid_sat_names"]) - if model.baseline_mode: - embeddings = model(drone_img=drone_img, sat_img=sat_img) - else: - embeddings = model( - drone_img=drone_img, - sat_img=sat_img, - caption_l1=batch["caption_l1"], - caption_l2=batch["caption_l2"], - caption_l3=batch["caption_l3"], - sat_caption_l1=batch["sat_caption_l1"], - sat_caption_l2=batch["sat_caption_l2"], - sat_caption_l3=batch["sat_caption_l3"], - ) - all_query.append(embeddings["query"].cpu()) - all_gallery.append(embeddings["gallery"].cpu()) - all_sat_names.extend(batch["sat_names"]) - - # Per-batch loss. + # Per-batch loss: use first valid sat per query as its paired gallery. if loss_fn is not None: - loss_dict = loss_fn(embeddings, epoch=epoch, total_epochs=total_epochs) - batch_losses.append(float(loss_dict["total"].item())) + pair_indices: list[int] = [] + for names in batch["valid_sat_names"]: + for name in names: + if name in sat_name_to_idx: + pair_indices.append(sat_name_to_idx[name]) + break + else: + pair_indices.append(-1) + if all(idx >= 0 for idx in pair_indices): + paired_gallery = gallery[pair_indices].to(device) + fake_embeddings = { + "query": q, + "gallery": paired_gallery, + "gate_q": model.fusion_query.gate_value, + "gate_g": model.fusion_gallery.gate_value, + } + loss_dict = loss_fn(fake_embeddings, epoch=epoch, total_epochs=total_epochs) + batch_losses.append(float(loss_dict["total"].item())) - query = torch.cat(all_query, dim=0) - gallery = torch.cat(all_gallery, dim=0) + query = torch.cat(query_embs, dim=0) # [N_q, D] + n_query = query.size(0) - sim = query @ gallery.t() - n = sim.size(0) + # --- Similarity + rankings --- + sim = query @ gallery.t() # [N_q, N_sat] + sorted_idx = sim.argsort(dim=1, descending=True) metrics: dict[str, float] = {} if batch_losses: metrics["loss"] = sum(batch_losses) / len(batch_losses) - # Build valid matches: for each query i, which gallery indices are valid? - # Get all valid sat names per query from the dataset. - dataset = loader.dataset - n_eval = min(n, len(dataset)) - if hasattr(dataset, "get_all_valid_sat_names"): - all_valid_names = dataset.get_all_valid_sat_names()[:n_eval] - else: - all_valid_names = None - - # Build sat_name → gallery index mapping. - sat_name_to_idx: dict[str, list[int]] = {} - for idx, name in enumerate(all_sat_names): - sat_name_to_idx.setdefault(name, []).append(idx) - - sorted_idx = sim.argsort(dim=1, descending=True) + # Precompute valid gallery index sets per query. + valid_idx_per_query: list[set[int]] = [] + for names in query_valid_names: + valid = {sat_name_to_idx[n] for n in names if n in sat_name_to_idx} + valid_idx_per_query.append(valid) # R@K with multi-match. for k in k_values: hits = 0 - for i in range(n_eval): - top_k_indices = sorted_idx[i, :k].tolist() - if all_valid_names is not None: - # Check if any valid satellite name appears in top-K gallery. - valid_gallery_indices = set() - for vname in all_valid_names[i]: - valid_gallery_indices.update(sat_name_to_idx.get(vname, [])) - if valid_gallery_indices.intersection(top_k_indices): - hits += 1 - else: - # Fallback: diagonal matching. - if i in top_k_indices: - hits += 1 - metrics[f"r@{k}_q2g"] = hits / max(n_eval, 1) + for i in range(n_query): + top_k = set(sorted_idx[i, :k].tolist()) + if valid_idx_per_query[i] & top_k: + hits += 1 + metrics[f"r@{k}_q2g"] = hits / max(n_query, 1) - # AP (mean reciprocal rank over valid matches). - ap_sum = 0.0 - for i in range(n_eval): - ranking = sorted_idx[i].tolist() - if all_valid_names is not None: - valid_gallery_indices = set() - for vname in all_valid_names[i]: - valid_gallery_indices.update(sat_name_to_idx.get(vname, [])) - # Find first valid match rank. - for rank, gidx in enumerate(ranking): - if gidx in valid_gallery_indices: - ap_sum += 1.0 / (rank + 1) - break - else: - for rank, gidx in enumerate(ranking): - if gidx == i: - ap_sum += 1.0 / (rank + 1) - break - metrics["ap_q2g"] = ap_sum / max(n_eval, 1) + # MRR over valid matches (kept key `ap_q2g` for CSV/plot compatibility). + mrr_sum = 0.0 + n_scored = 0 + for i in range(n_query): + valid = valid_idx_per_query[i] + if not valid: + continue + n_scored += 1 + for rank, gidx in enumerate(sorted_idx[i].tolist()): + if gidx in valid: + mrr_sum += 1.0 / (rank + 1) + break + metrics["ap_q2g"] = mrr_sum / max(n_scored, 1) + metrics["n_query"] = float(n_query) + metrics["n_gallery"] = float(gallery.size(0)) metrics["gate_q"] = model.fusion_query.gate_value metrics["gate_g"] = model.fusion_gallery.gate_value @@ -470,16 +516,31 @@ def train(cfg: TrainConfigGTAUAV) -> None: if tracker.has_wandb: tracker.watch_model(model, log_freq=50) - # Loss — WeightedInfoNCE for GTA-UAV (handles partial satellite overlap). - loss_fn = WeightedInfoNCELoss( - temperature_init=cfg.tau_init, - learnable_temperature=cfg.learnable_temperature, - label_smoothing=cfg.label_smoothing, - ) + # Loss. + if cfg.loss_type == "symmetric": + loss_fn = InfoNCELoss( + temperature_init=cfg.tau_init, + learnable_temperature=cfg.learnable_temperature, + label_smoothing=cfg.label_smoothing, + weight_q2g=cfg.weight_q2g, + weight_g2q=cfg.weight_g2q, + ) + loss_name = "SymmetricInfoNCE" + elif cfg.loss_type == "weighted": + loss_fn = WeightedInfoNCELoss( + temperature_init=cfg.tau_init, + learnable_temperature=cfg.learnable_temperature, + label_smoothing=cfg.label_smoothing, + ) + loss_name = "WeightedInfoNCE" + else: + raise ValueError(f"Unknown loss_type={cfg.loss_type!r} (expected 'symmetric' or 'weighted')") + LOGGER.info( - "Loss: WeightedInfoNCE Temperature: %s (init=%.3f)", + "Loss: %s Temperature: %s (init=%.3f) q2g=%.2f g2q=%.2f", + loss_name, "learnable" if cfg.learnable_temperature else "fixed", - cfg.tau_init, + cfg.tau_init, cfg.weight_q2g, cfg.weight_g2q, ) # Hard negative memory bank. @@ -509,15 +570,34 @@ def train(cfg: TrainConfigGTAUAV) -> None: filter_meta=cfg.filter_meta, ) - train_loader = DataLoader( - train_ds, - batch_size=cfg.batch_size, - shuffle=True, - num_workers=cfg.num_workers, - collate_fn=collate_gtauav_batch, - pin_memory=True, - drop_last=True, - ) + 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, + ) + LOGGER.info( + "Sampler: MutuallyExclusive — no false negatives within a batch", + ) + train_loader = DataLoader( + train_ds, + batch_sampler=mutex_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, + shuffle=True, + num_workers=cfg.num_workers, + collate_fn=collate_gtauav_batch, + pin_memory=True, + drop_last=True, + ) test_loader = DataLoader( test_ds, batch_size=cfg.batch_size, @@ -607,6 +687,8 @@ 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) epoch_start = time.time() agg: dict[str, float] = {} n_batches = 0 @@ -763,6 +845,8 @@ def train(cfg: TrainConfigGTAUAV) -> None: train_row["ap_q2g"] = train_recall.get("ap_q2g", 0.0) csv_logger.log_train(epoch, train_row, optimizer.param_groups[0]["lr"], elapsed) generate_plots(csv_logger.log_dir) + + if train_recall: LOGGER.info( "train-recall epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f", epoch,