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

24
CLAUDE.md
View File

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

23
conf/gtauav_balanced.gin
View File

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

46
scripts/smoke_eval.py Normal file
View File

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

79
scripts/smoke_train.py Normal file
View File

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

106
src/datasets/gtauav_dataset.py
View File

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

112
src/datasets/mutually_exclusive_sampler.py Normal file
View File

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

4
src/losses/multi_infonce.py
View File

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

2
src/losses/weighted_infonce.py
View File

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

99
src/models/asymmetric_encoder.py
View File

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

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,