When stripnet_freeze=False, all StripNet backbone params train end-to-end
with a separate optimizer group at lr * stripnet_backbone_lr_factor (default
0.1, so 1e-5 with default learning_rate=1e-4) — typical fine-tuning practice
for ImageNet-pretrained CNNs to avoid catastrophic forgetting.
Conv-MONA is now optional (stripnet_mona_last_n_stages=0 disables it). Three
modes are now supported:
- frozen + MONA: PEFT-style (~1.2M trainable, original default)
- unfrozen, no MONA: full fine-tune (~13.85M, all backbone params)
- unfrozen + MONA: hybrid (~14.5M, backbone + extra adapters)
_build_param_groups: new "backbone" group identifies image_encoder.backbone.*
params (excluding mona_*) when backbone="stripnet"; assigned lr factor
controls fine-tune step size independently from text/MONA groups.
conf/gtauav_balanced_stripnet_unfrozen.gin + baseline variant: ready-to-use
configs for full fine-tune experiment.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- _evaluate: compute R@K + AP for both directions (q2g and g2q) via inverted
ground truth; g2q denominator counts only sat-tiles with at least one positive
drone in the (sub)sampled query set. Surfaces in train.csv, val.csv,
train_recall.csv, W&B summary, and final log.
- conf/gtauav_balanced_asym.gin: asymmetric WEB+SAT encoders, MONA in all 24
ViT blocks (~17.6M trainable / ~733M total).
- conf/gtauav_baseline_asym.gin: same architecture, baseline_mode=True for
Δ R@1 against balanced_asym.
- CLAUDE.md / README.md: document new configs, clarify that g2q is now
computed (was claimed but missing).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Hard mining (c307269) uses `queue.t()` as a view directly in the
similarity matmul, so the autograd graph now holds a reference to the
memory-bank buffer's storage. The previous code enqueued fresh gallery
embeddings BEFORE `backward()`, which mutates that same buffer in place
and triggers:
RuntimeError: one of the variables needed for gradient computation has
been modified by an inplace operation: [torch.cuda.FloatTensor [1024, 8]]
is at version 2; expected version 1 instead.
The older InfoNCE path called `torch.cat([emb_b, queue], 0)`, which
materialises a fresh contiguous tensor and severs the alias — so the
same enqueue order worked. After the mining refactor, we need the
buffer to stay stable until backward completes.
Moving the enqueue to after backward is semantically identical: the
queue state observed by the next training step is the same either way
(FIFO, not involved in the just-completed backward).
`smoke_train.py` already had the correct order and didn't catch the
regression.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
InfoNCELoss.forward() doesn't accept `positive_weights` — that kwarg is
specific to WeightedInfoNCELoss's adaptive label-smoothing path. After
switching the default `loss_type` to "symmetric", training crashed with
`TypeError: unexpected keyword argument 'positive_weights'`.
Build the kwargs dict conditionally: add `positive_weights` only when
the loss is an instance of WeightedInfoNCELoss.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Three related architecture changes, driven by a cost/simplicity trade-off:
1. **Shared encoder**: one DINOv3 LVD-1689M (WEB) processes both drone
and satellite images. Previously asymmetric — separate WEB (drone) and
SAT-493M (satellite) encoders. Saves ~303M frozen params and halves
VRAM for the image tower. Expected to lose some satellite-domain
inductive bias; MONA adapters pick up the slack.
2. **MONA in last 12/24 blocks**: adapters injected only in the top half
of the ViT. The lowest 12 blocks keep their pretrained features
untouched. Trainable MONA count drops from 14.0M (48 adapters × 2
encoders) to 3.5M (24 adapters × 1 encoder).
3. **No DINO_SAT**: `nn_models/DINO_SAT` is no longer loaded by the
default config. It stays on disk and the path param is kept for
backward compat with asymmetric checkpoints.
Parameter counts (with text fusion + LoRA + gates):
Before: 17.6M trainable / 733M total (2.35%)
After: 7.06M trainable / 434M total (1.63%)
Also fixes a pre-existing resume bug: checkpoints now record
`shared_encoder`, `baseline_mode`, `mona_bottleneck`, `mona_last_n_blocks`
so `AsymmetricEncoder.load_checkpoint` can rebuild the right architecture.
Old checkpoints still load (missing keys fall back to asymmetric defaults
via `ckpt.get(..., <default>)`).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Three upgrades to the DynamicSimilaritySampler infrastructure:
1. **GPU kNN** (`dss_knn_device="cuda"`, default):
Moves the per-seed similarity matmul to the GPU. At 25K train items
this cuts per-epoch sampler generation from 17s to 1.6s — a 10.8x
speedup. Negligible VRAM (100MB for the [N, 1024] embedding tensor).
2. **LSH index** (`src/datasets/lsh_index.py`, opt-in via `dss_use_lsh=True`):
Random-projection cosine-LSH with H tables of B bits each. When enabled,
the sampler narrows the candidate pool per seed via hash-bucket lookup
before exact refinement. At 25K it's a wash (pool already fits in VRAM)
but provides a scaling path for 100K+ where the N² similarity matrix
would stop fitting. Default off.
3. **Embedding cache** (`src/datasets/embedding_cache.py`, `dss_cache_dir` config):
Disk-backed cache for drone query embeddings, keyed by epoch. Skips
re-embedding on --resume and lets ablations replay from a snapshot.
Atomic writes via `.tmp` → `.replace`.
Measured on 25K train entries, 1024-dim random embeddings:
CPU kNN: 17.44s
GPU kNN: 1.62s (10.8x)
GPU + LSH: 1.42s (LSH candidate pool 0.05% of N)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Batches assembled from visually-similar drone queries pressure the model to
learn finer discriminative features. Random mutex batches average ~0.26
pairwise cosine similarity in query embedding space; DSS batches average
~0.71 — confirming the lookalikes grouping works as intended.
Algorithm per batch:
1. Pick a random seed drone from the remaining pool.
2. Rank the entire remaining pool by cosine similarity to the seed.
3. Walk the ranking in descending order; add items whose sat_candidates
don't collide with the batch's already-claimed set.
4. Drop the seed if no valid batch can be assembled (rare mutex deadlock).
Inherits MutuallyExclusiveSampler semantics — no false negatives. Degrades
gracefully to mutex-only when no embeddings are set (warmup epochs, or if
`sampler_type="mutex"` is chosen).
Integration in `train_gtauav.py`:
- New `_embed_drone_queries` helper: model.encode_query forwarded over
GTAUAVDroneQuery, returns [N, D] CPU tensor. ~13s per 1024 queries on
a 4090 → ~5 min for the full 25K train set.
- Epoch loop re-embeds every `dss_reembed_every` epochs after a `dss_warmup_epochs`
warmup (first epochs use mutex-only since untrained embeddings aren't
informative for kNN).
- Config: `sampler_type` ∈ {"mutex", "dss"}. Default flipped to "dss".
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
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>
train.csv now includes eval_loss, r@1_q2g, r@5_q2g, r@10_q2g, ap_q2g
alongside training loss/temperature/gates when eval runs that epoch.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
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>
Architecture changes:
- Asymmetric DINOv3: WEB (drone) + SAT (satellite) with separate MONA
- MONA on all 24 blocks per encoder (was last 12)
- Remove projection, native 1024-dim retrieval space (was 512)
- Total: 748M params, 17.6M trainable (2.35%)
Hard negative memory bank:
- MoCo-style FIFO queue of 4096 detached gallery embeddings
- Each batch: B in-batch + Q queue negatives in InfoNCE
- Queue updated after each forward pass
Training config:
- batch_size=8, grad_accum=8, effective_batch=64
- eval_every=1 (eval + train recall every epoch)
- Max bs=24 with grad checkpointing on RTX 4090
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Compute AP (Mean Reciprocal Rank) in _evaluate() for both q→g and g→q
- AP saved in val.csv and train_recall.csv alongside R@K
- New AP plot panel in val_metrics.png (train vs val, both directions)
- Log AP in console output for train-recall and val epochs
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- _evaluate() now computes per-batch loss when loss_fn is provided
- Val loss and train recall loss saved in val.csv and train_recall.csv
- Overview plot shows train vs val loss curves side by side
- Helps detect overfitting: val loss diverging from train loss
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Evaluate R@K on train set (subset matching test size) alongside val
- New train_recall.csv with per-epoch train R@1/R@5/R@10
- Plot train vs val recall on same chart (solid=val, dashed=train)
- Helps detect overfitting: train R@1 up + val R@1 flat = overfit
- Train eval uses clean transforms (no augmentation)
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
epoch_000_train.csv and epoch_000_val.csv removed — all data is in
train.csv and val.csv which persist across resume.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
On resume, the same epoch may be re-run. Now log_train/log_val remove
existing entries for the same epoch before appending, preventing
duplicate rows in train.csv/val.csv.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
On init, load existing train.csv/val.csv so that epoch-level metrics
and plots include the full training history after checkpoint resume.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- MONA fp16 causes NaN (gamma=1e-6 underflows in fp16 min subnormal ~6e-8)
- Revert MONA forward to fp32 with autocast(enabled=False), cast output back
- Fix loss CSV: save raw_loss before backward() (tensor consumed after backward)
- Verified: loss=3.78, no NaN, bs=48 peak=21.4 GB
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- DINOv3: checkpoint each of 24 transformer blocks (recompute on backward)
- DGTRS-CLIP: checkpoint each of 12 transformer blocks
- Enables batch_size=24 on RTX 4090 (was 8 without checkpointing)
- Peak VRAM: 20.3 GB at bs=24 (was OOM at bs=16 before)
- ~20-30% slower per step, but 3x more in-batch negatives (23 vs 7)
- Enabled by default (gradient_checkpointing=True in config)
- Update README with VRAM benchmarks and checkpointing docs
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Single DINOv3 WEB for both drone and satellite branches (shared_encoder=True default)
- One set of MONA adapters instead of two: 7M trainable vs 14M
- Total params: 438M (was 748M), trainable: 10.6M (was 17.6M)
- Asymmetric mode still available via shared_encoder=False
- Add gradient accumulation (grad_accum_steps, --grad-accum CLI flag)
- Update model summary in README
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- New config field grad_accum_steps (default=1, no change in behavior)
- Loss scaled by 1/accum, optimizer step every N micro-batches
- Scheduler counts optimizer steps (not micro-batches)
- CLI flag --grad-accum for override
- Document gradient accumulation and in-batch negatives in README
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
train_batches.csv and epoch_N_batches.csv now update after every batch
instead of flushing at epoch end. Uses file append mode for efficiency.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Use gin.configurable(module=...) to prevent __main__ vs module name clash
- Remove `import src.training.train_gtauav` from gin files (already loaded)
- Use short selector names (TrainConfigGTAUAV) in all gin configs
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Add unified experiment tracker (W&B + TensorBoard) with graceful fallback
- Add gradient norm monitoring per param group (MONA, LoRA, MLP, gates, tau)
- Add Grad-CAM visualization for DINOv3 drone/satellite encoders
- Add PyTorch Profiler wrapper + torchinfo model summary
- Add gin-config support to train_gtauav.py with CLI overrides
- Add v3 gin configs: gtauav_balanced, gtauav_baseline, gtauav_text_heavy, gtauav_image_heavy
- Generate metric plots every epoch (not just on eval)
- Set default epochs to 10
- Update README and CLAUDE.md with new tooling and usage docs
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
New: src/training/plot_metrics.py
- train_metrics.png: loss, temperature, gates, lr
- val_metrics.png: R@K q→g and g→q
- overview.png: combined loss + R@1 + gates/tau
Auto-generates plots in {output_dir}/logs/ after each validation epoch.
Also callable standalone: python -m src.training.plot_metrics --log-dir ...
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Architecture changes:
- Removed proj_drone/proj_sat (1024→512): retrieval space is now
DINOv3 native 1024-dim, no information loss from projection
- TextFusionMLP: 2304→1024→1024 (was 2304→768→512), shared between branches
- Gallery branch now uses satellite captions (L1/L2/L3) via shared TextFusionMLP
- Two separate GatedFusion gates: α_q (query) and α_g (gallery)
- For sat images without captions (~57%): gate passes image features through
Dataset changes:
- GTAUAVDataset now loads satellite captions from caption index
- collate_gtauav_batch includes sat_caption_l1/l2/l3
Training loop:
- Passes satellite captions to model forward
- Logs both gate_q and gate_g values
11.1M trainable / 734M total (1.51%)
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Root cause: GradScaler scales gradients by ~65536 in fp16, causing
logit_scale.exp() gradient to overflow. The learnable temperature
and similarity logits must stay in fp32.
Fix: model forward runs inside autocast(fp16), but loss computation
(similarity @ temperature + cross_entropy) runs outside in fp32.
Also: clamp logit_scale in logit-space before exp() and force
similarity computation to fp32.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- scripts/make_split.py: merges cross-area train+test (33,708 pairs),
shuffles with seed=42, splits 80/20
- meta/train_80.json (26,966) + meta/test_20.json (6,742)
- After seg filter: 24,891 train / 6,252 test
- Default paths in train_gtauav.py updated to use new split
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- AsymmetricEncoder.save_checkpoint(): saves model_state + metadata
- AsymmetricEncoder.load_checkpoint(): rebuilds model with frozen backbones,
then loads trainable weights from checkpoint
- --resume flag restores optimizer, loss (learnable tau), and scheduler state
- Training continues from the saved epoch
Usage:
python -m src.training.train_gtauav --resume out/gtauav/with_text/ckpt_epoch004.pt
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
