Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix GTA-UAV evaluation and loss (critical: false negatives + wrong R@K)

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-24 12:40:10 +03:00
parent 04d5307221
commit b6dccbba7b
4 changed files with 242 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
conf/gtauav_balanced.gin
View File

@@ -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

22
src/datasets/gtauav_dataset.py
View File

@@ -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),
}

148
src/losses/weighted_infonce.py Normal file
View File

@@ -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(),
}

101
src/training/train_gtauav.py
View File

@@ -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,
)