Fix GTA-UAV evaluation and loss (critical: false negatives + wrong R@K)
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
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user