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

Fix GTA-UAV eval + training pipeline: full gallery, mutex sampler, per-sample mask

Browse Source 
Six critical fixes to the caption-test training/eval stack:

1. **IndentationError blocker** (train_gtauav.py:765-766)
   Unparseable file — train-recall LOGGER.info block was orphaned outside
   its `if eval_every` guard. Wrapped in `if train_recall:` so val eval
   and Grad-CAM only run on eval epochs.

2. **Full satellite gallery in `_evaluate`**
   Old code assembled gallery from DataLoader batches (one random sat per
   drone), producing an incomplete gallery of size ≈ N_query instead of
   N_unique_sat. Metrics were inflated because retrieval was against a
   subset that always contained the target.
   New `GTAUAVSatGallery` / `GTAUAVDroneQuery` iterate all unique tiles
   and queries independently; full-gallery multi-match R@K + MRR.

3. **Per-sample caption mask** (`AsymmetricEncoder._fuse_with_mask`)
   Mixed batches (some samples have captions, some don't) previously
   encoded empty strings through DGTRS and mixed the noise output into
   every sample via scalar gate. New `encode_query`/`encode_gallery` use
   `torch.where` to fall back to pure image features for empty-caption
   samples. Training `forward()` routes through the same helper so
   training and eval share code.

4. **Symmetric InfoNCE as primary loss** (multi_infonce.InfoNCELoss)
   Switched gin default from `WeightedInfoNCELoss` (adaptive label
   smoothing — not the Game4Loc soft-IoU target it claimed) to the
   existing symmetric InfoNCE with q2g=0.6/g2q=0.4 weighting. Loss type
   now selectable via `cfg.loss_type ∈ {"symmetric", "weighted"}`.

5. **MutuallyExclusiveSampler** (new file)
   BatchSampler that greedily packs drones whose `sat_candidates` sets
   are pairwise disjoint within a batch. Eliminates false negatives from
   the semi-positive graph without needing soft-label losses.
   At bs=8 keeps 100% of 24,891 train entries; at bs=64 keeps 92.6%.
   `set_epoch()` for reproducibility + different batches per epoch.

6. **Temperature clamp [0.01, 0.1]** (both loss modules)
   Old tau_max=0.5 allowed the logit distribution to collapse into a
   near-uniform softmax. Tightened to the CLIP-standard range.

Also:
- Added `scripts/smoke_eval.py` / `scripts/smoke_train.py` for fast
  regression checks (eval in ~2 min, 2 train steps in ~1 min on RTX 4090).
- CLAUDE.md updated to reflect the new pipeline.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-24 15:58:27 +03:00
parent ce7892926f
commit a499fcfd65
10 changed files with 640 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

286
src/training/train_gtauav.py
View File

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