From b6dccbba7bbf0b01f0505e841f93b1193f5897d8 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Fri, 24 Apr 2026 12:40:10 +0300 Subject: [PATCH] Fix GTA-UAV evaluation and loss (critical: false negatives + wrong R@K) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PROBLEM: GTA-UAV has overlapping satellite crops (partial IoU). Standard InfoNCE with diagonal targets treated valid matches as negatives. R@K checked only diagonal — missed valid matches, artificially low recall. FIXES: 1. WeightedInfoNCE loss (src/losses/weighted_infonce.py): - Per-sample adaptive label smoothing from positive_weights (IoU) - Higher weight → sharper target, lower → softer (semi-positive tolerance) - Based on Game4Loc reference implementation 2. Multi-match R@K evaluation: - Uses dataset.get_all_valid_sat_names() to get ALL valid matches per query - R@K counts hit if ANY valid satellite is in top-K (not just diagonal) - AP computed as MRR over first valid match 3. Dataset returns positive_weight per sample: - Sampled satellite weight passed to loss for adaptive smoothing - All valid satellite candidates exposed for evaluation Co-Authored-By: Claude Opus 4.6 (1M context) --- conf/gtauav_balanced.gin | 21 ++--- src/datasets/gtauav_dataset.py | 22 ++++- src/losses/weighted_infonce.py | 148 +++++++++++++++++++++++++++++++++ src/training/train_gtauav.py | 101 ++++++++++++++-------- 4 files changed, 242 insertions(+), 50 deletions(-) create mode 100644 src/losses/weighted_infonce.py diff --git a/conf/gtauav_balanced.gin b/conf/gtauav_balanced.gin index a958a2c..7a5247d 100644 --- a/conf/gtauav_balanced.gin +++ b/conf/gtauav_balanced.gin @@ -1,11 +1,8 @@ # GTA-UAV Balanced: Asymmetric DINOv3 (WEB+SAT) with L1/L2/L3 captions. -# query = sigma(alpha) * drone + (1-sigma(alpha)) * text -> InfoNCE vs gallery +# WeightedInfoNCE loss for GTA-UAV partial overlap handling. # 10 epochs, MONA all 24 blocks, 1024-dim retrieval, hard negative bank. -# -# NOTE: TrainConfigGTAUAV is registered by train_gtauav.py before gin parsing. -# InfoNCELoss is registered via import below. -import src.losses.multi_infonce +import src.losses.weighted_infonce # ---- Training ---- TrainConfigGTAUAV.epochs = 10 @@ -31,8 +28,6 @@ TrainConfigGTAUAV.gradient_checkpointing = True # ---- Loss ---- TrainConfigGTAUAV.tau_init = 0.07 TrainConfigGTAUAV.label_smoothing = 0.1 -TrainConfigGTAUAV.weight_q2g = 0.6 -TrainConfigGTAUAV.weight_g2q = 0.4 TrainConfigGTAUAV.learnable_temperature = True TrainConfigGTAUAV.neg_bank_size = 4096 @@ -47,10 +42,8 @@ TrainConfigGTAUAV.gradcam_every = 5 TrainConfigGTAUAV.use_profiler = False TrainConfigGTAUAV.log_grad_norms = True -# ---- InfoNCE Loss (gin-configurable) ---- -InfoNCELoss.temperature_init = 0.07 -InfoNCELoss.temperature_final = 0.01 -InfoNCELoss.label_smoothing = 0.1 -InfoNCELoss.weight_q2g = 0.6 -InfoNCELoss.weight_g2q = 0.4 -InfoNCELoss.learnable_temperature = True +# ---- WeightedInfoNCE (gin-configurable) ---- +WeightedInfoNCELoss.temperature_init = 0.07 +WeightedInfoNCELoss.learnable_temperature = True +WeightedInfoNCELoss.label_smoothing = 0.1 +WeightedInfoNCELoss.k = 5.0 diff --git a/src/datasets/gtauav_dataset.py b/src/datasets/gtauav_dataset.py index 5d8c1dc..3a0f02a 100644 --- a/src/datasets/gtauav_dataset.py +++ b/src/datasets/gtauav_dataset.py @@ -210,6 +210,14 @@ class GTAUAVDataset(Dataset): return transform(rgb) return torch.tensor(0) # placeholder if no transform + def get_all_valid_sat_names(self) -> list[list[str]]: + """Return all valid satellite matches per drone query (for evaluation). + + In GTA-UAV, each drone has multiple valid satellite crops (partial IoU). + Standard diagonal R@K is wrong — must check if ANY valid match is in top-K. + """ + return [entry["sat_candidates"] for entry in self.entries] + def __len__(self) -> int: return len(self.entries) @@ -220,13 +228,17 @@ class GTAUAVDataset(Dataset): # Sample satellite match (weighted if semi-positive). if entry["sat_weights"] is not None: - sat_name = self._rng.choices( - entry["sat_candidates"], + sat_idx = self._rng.choices( + range(len(entry["sat_candidates"])), weights=entry["sat_weights"], k=1, )[0] + sat_name = entry["sat_candidates"][sat_idx] + pos_weight = entry["sat_weights"][sat_idx] else: - sat_name = self._rng.choice(entry["sat_candidates"]) + sat_idx = self._rng.randrange(len(entry["sat_candidates"])) + sat_name = entry["sat_candidates"][sat_idx] + pos_weight = 1.0 # strict positive sat_img = self._load_image(entry["sat_dir"], sat_name, self.sat_transform) @@ -253,6 +265,8 @@ class GTAUAVDataset(Dataset): "sat_caption_l2": sat_l2, "sat_caption_l3": sat_l3, "pair_id": entry["drone_name"], + "sat_name": sat_name, + "positive_weight": pos_weight, } @@ -270,4 +284,6 @@ def collate_gtauav_batch( "sat_caption_l2": [b["sat_caption_l2"] for b in batch], "sat_caption_l3": [b["sat_caption_l3"] for b in batch], "pair_ids": [b["pair_id"] for b in batch], + "sat_names": [b["sat_name"] for b in batch], + "positive_weights": torch.tensor([b["positive_weight"] for b in batch], dtype=torch.float32), } diff --git a/src/losses/weighted_infonce.py b/src/losses/weighted_infonce.py new file mode 100644 index 0000000..7148a04 --- /dev/null +++ b/src/losses/weighted_infonce.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +"""Weighted InfoNCE loss for GTA-UAV cross-view geo-localization. + +Adapted from Game4Loc (https://github.com/Yux1angJi/GTA-UAV). +Uses per-sample label smoothing based on positive_weights (IoU/distance) +to handle partial overlap between drone and satellite crops. + +Standard InfoNCE assumes strict 1-to-1 pairs and treats all non-diagonal +entries as negatives. In GTA-UAV, multiple satellite crops can validly +match one drone image (partial IoU overlap), causing false negatives. +WeightedInfoNCE softens this with adaptive label smoothing per sample. +""" + +import math + +import gin +import torch +import torch.nn as nn +import torch.nn.functional as F + + +@gin.configurable +class WeightedInfoNCELoss(nn.Module): + """Weighted InfoNCE with adaptive per-sample label smoothing. + + For each sample i, eps_i = 1 - (1 - base_smoothing) / (1 + exp(-k * w_i)) + where w_i is the positive weight (e.g. IoU with matched satellite crop). + Higher weight → lower eps → sharper target (strong positive). + Lower weight → higher eps → softer target (weak/semi-positive). + + Args: + temperature_init: Initial temperature (or learnable logit_scale). + learnable_temperature: If True, temperature is learnable (CLIP-style). + label_smoothing: Base label smoothing (used when no weights provided). + k: Sigmoid steepness for weight → eps mapping. + tau_min: Min clamp for learnable temperature. + tau_max: Max clamp for learnable temperature. + """ + + def __init__( + self, + temperature_init: float = 0.07, + learnable_temperature: bool = True, + label_smoothing: float = 0.1, + k: float = 5.0, + tau_min: float = 0.01, + tau_max: float = 0.5, + ) -> None: + super().__init__() + self.label_smoothing = label_smoothing + self.k = k + self.tau_min = tau_min + self.tau_max = tau_max + self.learnable_temperature = learnable_temperature + + if learnable_temperature: + self.logit_scale = nn.Parameter( + torch.tensor(math.log(1.0 / temperature_init)) + ) + else: + self.logit_scale = None + self.temperature = temperature_init + + @property + def current_temperature(self) -> float: + if self.logit_scale is not None: + tau = 1.0 / self.logit_scale.exp().clamp( + min=1.0 / self.tau_max, max=1.0 / self.tau_min, + ).item() + return tau + return self.temperature + + def _compute_eps(self, positive_weights: torch.Tensor | None, n: int) -> torch.Tensor | list[float]: + """Compute per-sample label smoothing from positive weights.""" + if positive_weights is not None: + # Higher weight → lower eps (sharper, stronger positive). + return 1.0 - (1.0 - self.label_smoothing) / (1.0 + torch.exp(-self.k * positive_weights)) + return [self.label_smoothing] * n + + def _weighted_loss( + self, + sim_matrix: torch.Tensor, + eps_all: torch.Tensor | list[float], + ) -> torch.Tensor: + """Weighted InfoNCE: per-sample interpolation between hard and uniform targets. + + For each row i: + L_i = (1-eps_i) * [-sim[i,i] + logsumexp(sim[i,:])] + + eps_i * [-mean(sim[i,:]) + logsumexp(sim[i,:])] + """ + n = sim_matrix.shape[0] + total_loss = torch.tensor(0.0, device=sim_matrix.device) + for i in range(n): + eps = eps_all[i] if isinstance(eps_all, list) else eps_all[i] + logsumexp = torch.logsumexp(sim_matrix[i, :], dim=0) + total_loss += (1 - eps) * (-sim_matrix[i, i] + logsumexp) + total_loss += eps * (-sim_matrix[i, :].mean() + logsumexp) + return total_loss / n + + def forward( + self, + embeddings: dict[str, torch.Tensor], + epoch: int = 0, + total_epochs: int = 1, + positive_weights: torch.Tensor | None = None, + queue_negatives: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Compute weighted InfoNCE loss. + + Args: + embeddings: Dict with 'query' [B,D], 'gallery' [B,D], 'gate_q', 'gate_g'. + positive_weights: Per-sample weight [B] (e.g. IoU with matched sat crop). + queue_negatives: Extra negatives [Q,D] from memory bank (not used with weighted loss). + """ + query = embeddings["query"].float() + gallery = embeddings["gallery"].float() + + # Temperature. + if self.learnable_temperature: + clamped = self.logit_scale.float().clamp( + min=math.log(1.0 / self.tau_max), + max=math.log(1.0 / self.tau_min), + ) + logit_scale = clamped.exp() + tau = 1.0 / logit_scale + else: + logit_scale = 1.0 / self.temperature + tau = self.temperature + + sim_q2g = logit_scale * query @ gallery.t() + sim_g2q = sim_q2g.t() + + eps = self._compute_eps(positive_weights, query.shape[0]) + + loss_q2g = self._weighted_loss(sim_q2g, eps) + loss_g2q = self._weighted_loss(sim_g2q, eps) + total = (loss_q2g + loss_g2q) / 2 + + gate_q = embeddings.get("gate_q", 1.0) + gate_g = embeddings.get("gate_g", 1.0) + + return { + "total": total, + "temperature": tau if isinstance(tau, torch.Tensor) else torch.tensor(tau, device=total.device), + "gate_q": torch.tensor(gate_q, device=total.device) if not isinstance(gate_q, torch.Tensor) else gate_q.detach(), + "gate_g": torch.tensor(gate_g, device=total.device) if not isinstance(gate_g, torch.Tensor) else gate_g.detach(), + } diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 230bf74..53a34f3 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -30,7 +30,7 @@ from torch.utils.data import DataLoader from tqdm import tqdm from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch -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 from src.training.trackers import ExperimentTracker @@ -97,8 +97,6 @@ class TrainConfigGTAUAV: # Loss. tau_init: float = 0.07 label_smoothing: float = 0.1 - weight_q2g: float = 0.6 - weight_g2q: float = 0.4 learnable_temperature: bool = True neg_bank_size: int = 4096 # hard negative memory bank size (0 = disabled) @@ -184,10 +182,16 @@ def _evaluate( max_batches: int | None = None, desc: str = "eval", ) -> dict[str, float]: - """Compute R@K and optional loss. Use max_batches to limit for train set.""" + """Compute R@K with multi-match support for GTA-UAV. + + 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). + """ model.eval() all_query: list[torch.Tensor] = [] all_gallery: list[torch.Tensor] = [] + all_sat_names: list[str] = [] batch_losses: list[float] = [] for i, batch in enumerate(tqdm(loader, desc=f" {desc}", unit="batch", leave=False)): @@ -211,8 +215,9 @@ def _evaluate( ) all_query.append(embeddings["query"].cpu()) all_gallery.append(embeddings["gallery"].cpu()) + all_sat_names.extend(batch["sat_names"]) - # Per-batch loss (if loss_fn provided). + # Per-batch loss. 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())) @@ -222,34 +227,64 @@ def _evaluate( sim = query @ gallery.t() n = sim.size(0) - targets = torch.arange(n) metrics: dict[str, float] = {} - - # Average loss across batches. if batch_losses: metrics["loss"] = sum(batch_losses) / len(batch_losses) - # R@K and AP (q→g). + # 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) + + # R@K with multi-match. for k in k_values: - top_k = sorted_idx[:, :k] - hit = (top_k == targets.unsqueeze(1)).any(dim=1).float() - metrics[f"r@{k}_q2g"] = float(hit.mean().item()) + 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) - # AP q→g: for each query, rank of the correct gallery = 1/(rank+1). - ranks_q2g = (sorted_idx == targets.unsqueeze(1)).nonzero(as_tuple=True)[1].float() - metrics["ap_q2g"] = float((1.0 / (ranks_q2g + 1)).mean().item()) - - # R@K and AP (g→q). - sorted_idx_g2q = sim.t().argsort(dim=1, descending=True) - for k in k_values: - top_k = sorted_idx_g2q[:, :k] - hit = (top_k == targets.unsqueeze(1)).any(dim=1).float() - metrics[f"r@{k}_g2q"] = float(hit.mean().item()) - - ranks_g2q = (sorted_idx_g2q == targets.unsqueeze(1)).nonzero(as_tuple=True)[1].float() - metrics["ap_g2q"] = float((1.0 / (ranks_g2q + 1)).mean().item()) + # 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) metrics["gate_q"] = model.fusion_query.gate_value metrics["gate_g"] = model.fusion_gallery.gate_value @@ -435,17 +470,15 @@ def train(cfg: TrainConfigGTAUAV) -> None: if tracker.has_wandb: tracker.watch_model(model, log_freq=50) - # Loss. - loss_fn = InfoNCELoss( + # Loss — WeightedInfoNCE for GTA-UAV (handles partial satellite overlap). + loss_fn = WeightedInfoNCELoss( temperature_init=cfg.tau_init, - label_smoothing=cfg.label_smoothing, - weight_q2g=cfg.weight_q2g, - weight_g2q=cfg.weight_g2q, learnable_temperature=cfg.learnable_temperature, + label_smoothing=cfg.label_smoothing, ) LOGGER.info( - "Temperature: %s (init=%.3f)", - "learnable" if cfg.learnable_temperature else "cosine schedule", + "Loss: WeightedInfoNCE Temperature: %s (init=%.3f)", + "learnable" if cfg.learnable_temperature else "fixed", cfg.tau_init, ) @@ -608,12 +641,14 @@ def train(cfg: TrainConfigGTAUAV) -> None: sat_caption_l2=batch["sat_caption_l2"], sat_caption_l3=batch["sat_caption_l3"], ) - # Loss in fp32 with optional hard negative queue. + # Loss — WeightedInfoNCE with positive weights from dataset. + pos_weights = batch["positive_weights"].to(cfg.device, non_blocking=True) queue_neg = neg_bank.get_queue() if neg_bank is not None else None loss_dict = loss_fn( embeddings=embeddings, epoch=epoch, total_epochs=cfg.epochs, + positive_weights=pos_weights, queue_negatives=queue_neg, )