Files
caption-test/README.md
pikaliov 814586ce3b Add g2q eval metrics and asymmetric MONA-24 configs
- _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>
2026-04-25 13:47:02 +03:00

587 lines
28 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# Caption Quality Test for Cross-View Geo-Localization
Validate whether generated text captions improve retrieval R@1 in cross-view
geo-localization (drone-to-satellite).
## Architecture
### Overview
Shared DINOv3 WEB encoder with MONA adapters (last 12 blocks, bfloat16),
DGTRS-CLIP text fusion, and projection to 512-dim retrieval space.
```
┌─────────────────────────────── QUERY BRANCH ─────────────────────────────────┐
│ │
│ drone_img ──► DINOv3 ViT-L/16 WEB (frozen, 303M) ──► CLS [B,1024] │
│ [B,3,256,256] + MONA adapters (3.5M, bf16) │ │
│ (2 per block × last 12 of 24, bn=64) │ │
│ + gradient checkpointing Projection │
│ Linear(1024→512) │
│ d_img [B,512] │
│ │ │
│ L1 (overview) ──► DGTRS-CLIP ViT-L-14 ──► z₁ [B,768] ─┐ │
│ L2 (full desc) ──► (frozen, 124M) ──► z₂ [B,768] ─┼─ cat ──[B,2304] │
│ L3 (fingerprint) ──► + LoRA r=4 (147K) ──► z₃ [B,768] ─┘ │ │
│ (248 tok, KPS pos.) TextFusionMLP (shared, 1.5M) │
│ + gradient checkpointing Linear(2304→512)→GELU→ │
│ Linear(512→512) │
│ │ │
│ d_txt [B,512] │
│ │ │
│ q = σ(α_q)·d_img + (1σ(α_q))·d_txt GatedFusion_q │
│ │ │
│ q̂ = q / ‖q‖₂ ──► query [B,512] │
└───────────────────────────────────────────────────────────────────────────────┘
┌────────────────────────────── GALLERY BRANCH ────────────────────────────────┐
│ │
│ sat_img ──► DINOv3 WEB (shared encoder) ──► CLS ──► Projection │
│ [B,3,256,256] + MONA (shared adapters) s_img [B,512] │
│ │ │
│ sat_L1 ──► DGTRS-CLIP ──► z₁ ─┐ │
│ sat_L2 ──► (shared) ──► z₂ ─┼─ cat ──► TextFusionMLP ──► s_txt [B,512] │
│ sat_L3 ──► + LoRA ──► z₃ ─┘ │ │
│ │ │
│ g = σ(α_g)·s_img + (1σ(α_g))·s_txt GatedFusion_g │
│ │ │
│ ĝ = g / ‖g‖₂ ──► gallery [B,512] │
└───────────────────────────────────────────────────────────────────────────────┘
┌────────────────────────────── RETRIEVAL ──────────────────────────────────────┐
│ │
│ similarity = q̂ · ĝᵀ / τ (τ learnable, init=0.07, clamp [0.01, 0.5]) │
│ loss = 0.6·CE(q→g) + 0.4·CE(g→q) (label_smoothing=0.1) │
│ │
│ Retrieval space: 512-dim (DINOv3 1024 → projection → 512) │
│ All shared: one DINOv3, one MONA set, one DGTRS-CLIP, one TextFusionMLP │
│ For sat images without captions: s_txt=None → g = s_img (gate passthrough) │
│ BASELINE: σ(α) = 1.0 for both branches (text disabled, DGTRS not loaded) │
└───────────────────────────────────────────────────────────────────────────────┘
┌──────────────────────────── DIAGNOSTICS PIPELINE ────────────────────────────┐
│ │
│ Per-batch ──► CSV (train_batches.csv) ──► TensorBoard / W&B scalars │
│ Per-epoch ──► CSV (train.csv, val.csv) ──► Seaborn plots (PNG) │
│ Every N epochs ──► Grad-CAM heatmaps (drone + satellite DINOv3 last block) │
│ Epoch 0 (opt) ──► PyTorch Profiler (Chrome trace + CUDA timeline) │
│ Per-50-batch ──► Gradient norms per group (MONA, LoRA, MLP, gates, τ) │
│ On init ──► torchinfo model summary (model_summary.txt) │
└───────────────────────────────────────────────────────────────────────────────┘
```
### Text hierarchy (L1 / L2 / L3)
Each drone image has a VLM-generated caption (Qwen3-VL) split into 3 levels:
| Level | Name | Content | Typical length |
|-------|------|---------|----------------|
| **L1** | Overview | First sentence of P1: land-cover summary with class percentages | 1530 tokens |
| **L2** | Full description | Complete P1 (inventory) + P2 (spatial layout with 5+ zones) | 100200 tokens |
| **L3** | Fingerprint | P3: unique landmarks and spatial signature for matching | 2050 tokens |
All three levels are encoded by a **single DGTRS-CLIP ViT-L-14** text encoder
(248-token context via KPS positional embedding, 768-dim output)
adapted with **LoRA** (rank=4 on Q/V in all 12 blocks).
### Text fusion (shared MLP for both branches)
```math
\mathbf{z}_{\text{text}} = \text{MLP}\bigl([\mathbf{z}_1 \;;\; \mathbf{z}_2 \;;\; \mathbf{z}_3]\bigr)
```
where `[z₁ ; z₂ ; z₃] ∈ ^{B×2304}` is the concatenation of three 768-dim DGTRS-CLIP embeddings, and
```math
\text{MLP}: \text{Linear}(2304, 512) \to \text{GELU} \to \text{Linear}(512, 512), \quad \mathbf{z}_{\text{text}} \in \mathbb{R}^{B \times 512}
```
### Gated fusion (separate gates for query and gallery)
```math
\mathbf{q} = \sigma(\alpha_q) \cdot \mathbf{d}_{\text{img}} + \bigl(1 - \sigma(\alpha_q)\bigr) \cdot \mathbf{d}_{\text{txt}} \qquad \text{(query branch)}
```
```math
\mathbf{g} = \sigma(\alpha_g) \cdot \mathbf{s}_{\text{img}} + \bigl(1 - \sigma(\alpha_g)\bigr) \cdot \mathbf{s}_{\text{txt}} \qquad \text{(gallery branch)}
```
- `α_q, α_g` — separate learnable scalars in logit-space, init `σ(α) ≈ 0.7`
- `σ` — sigmoid function
- `d_img, s_img ∈ ^{B×512}` — DINOv3+MONA → projection(1024→512)
- `d_txt, s_txt ∈ ^{B×512}` — fused text embeddings (TextFusionMLP → 512)
- For satellite images without captions: `s_txt = None → g = s_img`
### Adaptation methods
| Method | Applied to | Where | Params |
|--------|-----------|-------|--------|
| **MONA** (CVPR 2025) | DINOv3 ViT-L/16 (shared) | After MSA and MLP in each of 24 blocks | 6.85M |
| **LoRA** (rank=4) | DGTRS-CLIP text encoder | Q and V projections in all 12 blocks | 147K |
| **Projection** | After DINOv3 CLS output | Linear(1024→512) | 525K |
**MONA adapter** (per block):
```math
\mathbf{x} \leftarrow \mathbf{x} + \text{Up}_{64 \to 1024}\!\Bigl(\text{GELU}\bigl(\text{MonaOp}\bigl(\text{Down}_{1024 \to 64}(\hat{\mathbf{x}})\bigr)\bigr)\Bigr)
```
where `x̂ = γ · LN(x) + γₓ · x` (scaled LayerNorm, `γ` init `10⁻⁶`, `γₓ` init `1`)
```math
\text{MonaOp}(\mathbf{x}) = \frac{\text{DWConv}_{3 \times 3}(\mathbf{x}) + \text{DWConv}_{5 \times 5}(\mathbf{x}) + \text{DWConv}_{7 \times 7}(\mathbf{x})}{3} + \mathbf{x}
```
MONA runs in **bfloat16** with gradient checkpointing to save VRAM.
Applied only to the **last 12 blocks** (out of 24) — early blocks extract low-level
features (edges, textures) that are domain-agnostic and don't need spatial adaptation.
**Why bfloat16, not fp16:**
MONA's scaled LayerNorm uses `gamma` initialized at `1e-6` for near-identity output at start.
fp16 has min normal ~6.1e-5, so `1e-6` falls into the subnormal range where precision collapses,
causing NaN after a few blocks. bfloat16 has the same exponent range as fp32 (min ~1.2e-38),
so `1e-6` is safely representable. RTX 4090 supports bf16 natively with comparable throughput.
| Precision | `1e-6` representable | MONA stable | VRAM (bs=48) |
|-----------|:---:|:---:|:---:|
| fp32 | yes | yes | 21.4 GB |
| **bfloat16** | **yes** | **yes** | **21.8 GB** |
| fp16 | subnormal (lossy) | **NaN** | — |
**Why MONA over LoRA for DINOv3:**
MONA uses multi-scale depthwise convolutions (3×3, 5×5, 7×7) that provide **spatial inductive bias**
critical for cross-view geo-localization. Drone images (oblique, 100-600m altitude) and satellite
images (nadir) exhibit a strong **geometric domain gap** — the same building looks spatially different
from each viewpoint. MONA's multi-scale spatial filters learn scale-invariant features to bridge
this gap, while LoRA (pure linear low-rank correction) would only handle style/distribution shifts.
| | MONA | LoRA (on DINOv3) |
|---|---|---|
| Inductive bias | 2D spatial (knows about pixel neighbors) | None (linear correction) |
| Best for | Geometric domain gap (aerial↔satellite) | Style/distribution shift |
| Params | 6.85M (bottleneck=64) | ~0.3M (rank=4) |
| Compute | Heavier (192 conv ops per forward) | Light |
| CVGL fit | Strong (multi-scale spatial adaptation) | Weak (no spatial awareness) |
**LoRA** (per DGTRS-CLIP attention layer):
```math
\mathbf{Q}' = \mathbf{Q} + \frac{\alpha}{r} \cdot \mathbf{x} \mathbf{A}_Q^T \mathbf{B}_Q^T, \qquad \mathbf{V}' = \mathbf{V} + \frac{\alpha}{r} \cdot \mathbf{x} \mathbf{A}_V^T \mathbf{B}_V^T
```
where `A ∈ ^{r×d}`, `B ∈ ^{d×r}`, `r = 4`. LoRA is appropriate for the text encoder since
text has no spatial structure — the adaptation needed is purely semantic/distributional.
### Projection head (DINOv3 1024 → 512)
```math
\mathbf{h} = \text{Linear}_{1024 \to 512}\!\bigl(\text{CLS}_{\text{DINOv3}}\bigr)
```
Reduces the retrieval space from DINOv3 native 1024-dim to 512-dim. Benefits:
- Smaller similarity matrix in InfoNCE (`B×B` at 512 vs 1024)
- TextFusionMLP outputs 512 instead of 1024 (fewer params: 1.5M vs 3.4M)
- Shared projection for both drone and satellite branches
### Pair formation and negative sampling
The dataset provides only **positive pairs** — each entry maps one drone image to its
matching satellite crop(s). Negative pairs are **not** stored explicitly; instead, they
are constructed automatically inside each training batch via the InfoNCE loss:
```
Batch (B = 8):
drone_0 ↔ sat_0 ← positive (from dataset)
drone_1 ↔ sat_1 ← positive
...
drone_7 ↔ sat_7 ← positive
Similarity matrix B×B:
sat_0 sat_1 ... sat_7
q_0 [ pos neg ... neg ] ← CE target = 0
q_1 [ neg pos ... neg ] ← CE target = 1
...
q_7 [ neg neg ... pos ] ← CE target = 7
```
- **Positives:** diagonal `sim[i, i]` — correct drone-satellite pair
- **Negatives:** off-diagonal `sim[i, j≠i]` — other satellite images in the batch (in-batch negatives)
- At `batch_size=8`: 1 positive + 7 in-batch negatives per query
- No hard-negative mining — negatives are random within the batch
- Larger batch → more negatives → harder contrastive task
Each drone image may have multiple satellite candidates (with distance-based weights).
At training time, one satellite crop is sampled per drone (weighted if semi-positive).
### Loss function
Symmetric InfoNCE with learnable temperature (CLIP-style `logit_scale`):
```math
\mathcal{L} = w_{q \to g} \cdot \mathcal{L}_{q \to g} + w_{g \to q} \cdot \mathcal{L}_{g \to q}
```
```math
\mathcal{L}_{q \to g} = \text{CrossEntropy}\!\left(\frac{\hat{\mathbf{q}} \cdot \hat{\mathbf{g}}^T}{\tau},\; \text{targets}\right), \qquad \mathcal{L}_{g \to q} = \text{CrossEntropy}\!\left(\frac{\hat{\mathbf{g}} \cdot \hat{\mathbf{q}}^T}{\tau},\; \text{targets}\right)
```
- `τ = 1 / exp(logit_scale)` — learnable scalar, clamped `τ ∈ [0.01, 0.5]`, init `τ₀ = 0.07`
- `w_{q→g} = 0.6`, `w_{g→q} = 0.4`
- `targets = [0, 1, 2, ..., B1]` — positives on diagonal
- label smoothing `= 0.1`
Loss and adapters run in **fp32** (AMP autocast disabled) to prevent gradient overflow.
### Gradient checkpointing
Gradient checkpointing trades compute for VRAM by recomputing activations during
backward instead of storing them. Enabled by default (`gradient_checkpointing=True`).
| Component | Without checkpointing | With checkpointing |
|-----------|:---:|:---:|
| DINOv3 (24 blocks) | stores all 24 block activations | recomputes on backward |
| DGTRS-CLIP (12 blocks) | stores all 12 block activations | recomputes on backward |
| **Max batch_size** (RTX 4090, shared encoder) | **8** | **24** |
| Speed penalty | — | ~20-30% slower per step |
VRAM tested on RTX 4090 (24 GB) with shared DINOv3 WEB + MONA top-12 bf16 + DGTRS-CLIP + text:
| `batch_size` | Peak VRAM | Status |
|:---:|:---:|:---:|
| 32 | 16.1 GB | OK |
| 40 | 19.1 GB | OK |
| **48** | **21.8 GB** | **OK (recommended)** |
| 56 | >24 GB | OOM |
### Gradient accumulation
Gradient accumulation emulates a larger effective batch without extra memory:
```
effective_batch_size = batch_size × grad_accum_steps
```
| Setting | `batch_size` | `grad_accum_steps` | Effective batch | In-batch negatives |
|---------|:---:|:---:|:---:|:---:|
| Default | 48 | 1 | 48 | 47 |
| Large effective batch | 48 | 4 | 192 | 47 per micro-batch |
**Note:** gradient accumulation averages gradients across micro-batches, but each
micro-batch still only sees `batch_size` in-batch negatives. To increase the number
of negatives per forward pass, increase `batch_size` directly (requires more VRAM).
```bash
# Example: effective batch of 192 with gradient accumulation
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
--filter-meta meta/seg_filter.json --batch-size 48 --grad-accum 4
```
### Metrics
| Metric | Formula | Direction | Computed on |
|--------|---------|-----------|:-----------:|
| **R@K** (Recall at K) | fraction of queries where correct gallery is in top-K | q→g and g→q | train + val |
| **AP** (Average Precision) | mean of 1/(rank+1) across all queries (MRR) | q→g and g→q | train + val |
| **Loss** (InfoNCE) | symmetric cross-entropy on similarity matrix | — | train + val |
| **Delta R@1** | R@1(with_text) R@1(baseline) | q→g | val only |
R@1, R@5, R@10, AP, and loss are computed on both **train** (subset matching test size, clean
transforms) and **val** (full test set) every epoch. Train vs val comparison enables overfitting
detection: if train R@1/AP rises while val stagnates → overfitting.
**Output CSVs:**
| File | Content | Updated |
|------|---------|---------|
| `logs/train.csv` | Epoch-level train loss, temperature, gates, lr | Every epoch |
| `logs/val.csv` | Val R@K, AP, loss, gates | Every eval epoch |
| `logs/train_recall.csv` | Train R@K, AP, loss (subset) | Every eval epoch |
| `logs/train_batches.csv` | Per-batch loss, temperature, gates, lr | Every batch |
**Plots** (auto-generated in `logs/`):
| Plot | Panels |
|------|--------|
| `train_metrics.png` | Loss, temperature (τ), gates (σ(α)), learning rate |
| `val_metrics.png` | R@K q→g (train vs val), R@K g→q (train vs val), AP (train vs val) |
| `overview.png` | Train+val loss, val R@1, gates + temperature |
### Optimizer & scheduler
```
Optimizer: AdamW
- MONA adapters, projection, TextFusionMLP, gate α_q, gate α_g, logit_scale:
lr = 1e-4, weight_decay = 1e-4
- LoRA adapters (DGTRS-CLIP text encoder):
lr = 1e-5 (10× lower, --text-lr-factor 0.1)
Scheduler: Linear warmup (2 epochs) + cosine annealing
- Per optimizer step (accounts for gradient accumulation)
- warmup: lr linearly ramps from 0 to lr_max over warmup_steps
- cosine: lr decays from lr_max to 0 over remaining steps
Gradient clipping: max_norm = 1.0
Gradient accumulation: configurable (default 1, --grad-accum N)
Mixed precision: AMP fp16 for DINOv3/DGTRS forward, bf16 for MONA, fp32 for loss
Gradient checkpointing: DINOv3 (24 blocks) + DGTRS-CLIP (12 blocks)
```
### Augmentations
| Transform | Drone (train) | Satellite (train) | Eval |
|-----------|:---:|:---:|:---:|
| RandomResizedCrop(256, scale=0.71.0) | ✓ | ✓ | — |
| Resize(256) + CenterCrop(256) | — | — | ✓ |
| RandomHorizontalFlip(0.5) | ✓ | ✓ | — |
| RandomRotation(15°) | ✓ | — | — |
| ColorJitter(0.3, 0.3, 0.2, 0.1) | ✓ | ✓ | — |
| RandomGrayscale(0.05) | ✓ | ✓ | — |
| GaussianBlur(k=3, σ=0.12.0) | ✓ | — | — |
| ImageNet Normalize | ✓ | ✓ | ✓ |
### Model summary
| Component | Params | Trainable | Notes |
|-----------|--------|-----------|-------|
| DINOv3 ViT-L/16 WEB (shared) | 303M | frozen | single encoder for drone + satellite |
| MONA adapters (shared, bf16) | 3.5M | 3.5M | 2 per block × last 12 blocks, bottleneck=64 |
| Image projection | 525K | 525K | Linear(1024→512) after DINOv3 CLS |
| DGTRS-CLIP ViT-L-14 (text) | 124M | frozen | backbone weights frozen |
| LoRA adapters (text) | 147K | 147K | Q+V, rank=4, 12 blocks |
| TextFusionMLP (shared) | 1.5M | 1.5M | Linear(2304,512) + GELU + Linear(512,512) |
| GatedFusion α_q + α_g | 2 | 2 | separate gate scalars |
| logit_scale | 1 | 1 | learnable temperature |
| **Total (shared)** | **432M** | **5.6M (1.30%)** | retrieval dim = 512 |
> **Asymmetric mode** (`shared_encoder=False`): separate DINOv3 WEB (drone) + DINOv3 SAT
> (satellite) encoders with independent MONA adapters. Requires ~4-5 GB more VRAM.
> Use `conf/gtauav_balanced_asym.gin` / `conf/gtauav_baseline_asym.gin` — these set
> `shared_encoder=False` and `mona_last_n_blocks=24` for the full-capacity setup.
### Eval directions
`_evaluate` computes R@1/5/10 and AP (MRR) for both retrieval directions:
| Key | Direction | Notes |
|-----|-----------|-------|
| `r@K_q2g`, `ap_q2g` | drone → satellite (query → gallery) | Denominator: all queries (q2g convention) |
| `r@K_g2q`, `ap_g2q` | satellite → drone (gallery → query) | Denominator: only sat-tiles with ≥1 positive drone in subsample |
g2q is computed via inverted ground truth: for each sat-tile, collect the drone indices
that list it as a valid candidate. `n_scored_g2q` is reported in metrics for transparency.
## Experiments
### V3 — GTA-UAV + DINOv3 + DGTRS-CLIP (active)
**Dataset:** GTA-UAV-LR (33K drone + 14K satellite, GTA V synthetic)
- RGB: `/home/servml/Документы/datasets/GTA-UAV-LR/`
- Captions: `/home/servml/Документы/datasets/GTA-UAV-LR-captions/` (40K JSON, 3-paragraph VLM)
- Segmentation: `/home/servml/Документы/datasets/GTA-UAV-LR-aug/` (17 classes)
- Seg filter: 37,498 passed / 10,905 excluded (>=90% background+water)
- Split: 80/20 random (26,966 train / 6,742 test → 24,891/6,252 after seg filter)
### V2 — UAV-GeoLoc + GeoRSCLIP (legacy)
Single backbone with template captions.
```
Query: drone_img + caption -> GeoRSCLIP -> GatedFusion -> query
Gallery: sat_img -> GeoRSCLIP -> gallery
```
**Dataset:** UAV-GeoLoc Terrain split (206K train queries)
## Structure
```
caption-test/
├── conf/ # Gin configs
│ ├── gtauav_balanced.gin # GTA-UAV with text, shared encoder, MONA 12/24 (v3)
│ ├── gtauav_baseline.gin # GTA-UAV baseline, shared, MONA 12/24, no text (v3)
│ ├── gtauav_balanced_asym.gin # GTA-UAV with text, asymmetric WEB+SAT, MONA 24/24
│ ├── gtauav_baseline_asym.gin # GTA-UAV baseline, asymmetric, MONA 24/24, no text
│ ├── gtauav_text_heavy.gin # GTA-UAV text-heavy gate=0.3 (v3)
│ ├── gtauav_image_heavy.gin # GTA-UAV image-heavy gate=0.9 (v3)
│ ├── balanced.gin # UAV-GeoLoc with text (v2)
│ ├── baseline_no_text.gin # UAV-GeoLoc baseline (v2)
│ └── text_heavy.gin # UAV-GeoLoc text-heavy (v2)
├── nn_models/ # Pre-trained checkpoints (v3, gitignored)
│ ├── DINO_WEB/ # DINOv3 ViT-L/16 LVD-1689M (.pth)
│ ├── DINO_SAT/ # DINOv3 ViT-L/16 SAT-493M (.safetensors)
│ └── LRSCLIP/ # DGTRS-CLIP ViT-L-14 (.pt)
├── meta/ # Generated metadata
│ ├── train_80.json # 80% train split (26,966 pairs)
│ ├── test_20.json # 20% test split (6,742 pairs)
│ └── seg_filter.json # Segmentation filter results
├── scripts/
│ ├── make_split.py # Create 80/20 train/test split
│ ├── filter_segmentation.py # Exclude 90%+ background/water images
│ ├── compare_runs.py # Delta R@1 comparison report
│ └── generate_captions.py # Offline caption generation
├── src/
│ ├── datasets/
│ │ ├── gtauav_dataset.py # GTA-UAV-LR loader + L1/L2/L3 parsing (v3)
│ │ └── visloc_with_captions.py # UAV-GeoLoc loader (v2)
│ ├── models/
│ │ ├── dgtrs/ # Official DGTRS-CLIP text encoder (Apache-2.0)
│ │ │ ├── model.py # DGTRSTextEncoder, build_model, tokenize
│ │ │ ├── simple_tokenizer.py # BPE tokenizer (248 tokens)
│ │ │ └── bpe_simple_vocab_16e6.txt.gz
│ │ │ ├── adapters.py # MONA (bf16) + LoRA adapters
│ │ ├── asymmetric_encoder.py # DINOv3 + projection + DGTRS + GatedFusion (v3)
│ │ └── dual_encoder.py # GeoRSCLIP + GatedFusion (v2)
│ ├── losses/
│ │ └── multi_infonce.py # InfoNCE with learnable temperature
│ ├── training/
│ │ ├── train_gtauav.py # Training loop GTA-UAV (v3)
│ │ ├── train.py # Training loop UAV-GeoLoc (v2)
│ │ ├── trackers.py # Unified tracker: W&B + TensorBoard
│ │ ├── grad_monitor.py # Gradient norm monitoring per group
│ │ ├── gradcam.py # Grad-CAM visualization for DINOv3
│ │ ├── profiling.py # PyTorch Profiler + torchinfo summary
│ │ └── plot_metrics.py # Seaborn/matplotlib metric plots
│ └── eval/
│ └── evaluate.py # R@K metrics, Delta R@1
└── checkpoints/ # GeoRSCLIP RS5M_ViT-B-32.pt (v2)
```
## Prerequisites
```
torch>=2.0
safetensors
coloredlogs
tqdm
ftfy
regex
gin-config
Pillow
numpy
pandas
matplotlib
seaborn
```
### Optional (for extended diagnostics)
```
wandb # Weights & Biases experiment tracking
torchinfo # Model summary tables
tensorboard # TensorBoard logging (included with torch)
```
## Workflow (V3 — GTA-UAV)
### 1. Create 80/20 split and filter segmentation
```bash
python -m scripts.make_split --output-dir meta
python -m scripts.filter_segmentation --output meta/seg_filter.json
```
### 2. Train with gin configs (recommended)
```bash
# With captions (L1/L2/L3, bs=48, 30 epochs, eval every epoch)
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
--filter-meta meta/seg_filter.json --batch-size 48 --epochs 30 \
--gin-param 'TrainConfigGTAUAV.eval_every=1'
# Baseline (no text)
python -m src.training.train_gtauav --config conf/gtauav_baseline.gin \
--filter-meta meta/seg_filter.json --batch-size 48 --epochs 30
# Text-heavy (gate=0.3, 70% text weight)
python -m src.training.train_gtauav --config conf/gtauav_text_heavy.gin \
--filter-meta meta/seg_filter.json --batch-size 48 --epochs 30
# Image-heavy (gate=0.9, 10% text weight)
python -m src.training.train_gtauav --config conf/gtauav_image_heavy.gin \
--filter-meta meta/seg_filter.json --batch-size 48 --epochs 30
# Asymmetric variants: separate WEB (drone) + SAT (satellite) encoders, MONA in all 24 blocks
# Higher capacity (~733M total / ~17.6M trainable), larger VRAM footprint.
python -m src.training.train_gtauav --config conf/gtauav_balanced_asym.gin \
--filter-meta meta/seg_filter.json
python -m src.training.train_gtauav --config conf/gtauav_baseline_asym.gin \
--filter-meta meta/seg_filter.json
```
### 3. Train without gin (CLI-only)
```bash
python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.json --batch-size 48
python -m src.training.train_gtauav --filter-meta meta/seg_filter.json --batch-size 48
```
### 4. Enable diagnostics
```bash
# W&B + Grad-CAM + PyTorch Profiler
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
--filter-meta meta/seg_filter.json --batch-size 48 --wandb --gradcam --profile
# Gin parameter overrides from CLI
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
--filter-meta meta/seg_filter.json \
--gin-param 'TrainConfigGTAUAV.eval_every=1' 'TrainConfigGTAUAV.epochs=30'
```
CLI flags (`--wandb`, `--gradcam`, `--profile`, `--epochs`, `--batch-size`, etc.) take priority over gin config.
### 5. Resume from checkpoint
```bash
python -m src.training.train_gtauav --resume out/gtauav/with_text/ckpt_epoch004.pt \
--filter-meta meta/seg_filter.json
```
### 6. Compare and get verdict
```bash
python -m scripts.compare_runs \
--baseline_report out/gtauav/baseline/eval_report.json \
--full_report out/gtauav/with_text/eval_report.json \
--output out/gtauav/comparison.md
```
### 7. View TensorBoard
```bash
tensorboard --logdir out/gtauav/with_text/tb_logs
```
## Diagnostics & Visualization
| Tool | Flag | Output | Description |
|------|------|--------|-------------|
| **TensorBoard** | `--use-tb` (default on) | `{out}/tb_logs/` | Scalars, histograms, images |
| **W&B** | `--wandb` | cloud | Full experiment tracking, Grad-CAM images |
| **Grad-CAM** | `--gradcam` | `{out}/gradcam/` | DINOv3 attention heatmaps (drone + satellite) |
| **PyTorch Profiler** | `--profile` | `{out}/profiler/` | Chrome trace, CUDA timeline, memory |
| **torchinfo** | auto | `{out}/model_summary.txt` | Layer-by-layer parameter table |
| **Gradient norms** | `--log-grad-norms` (default on) | TB/W&B | Per-group: MONA, LoRA, MLP, gates, tau |
| **CSV (per-batch)** | auto | `{out}/logs/train_batches.csv` | Loss, tau, gates, lr for every batch |
| **CSV (per-epoch)** | auto | `{out}/logs/train.csv, val.csv` | Epoch loss averages + seaborn plots |
| **CSV (recall)** | auto | `{out}/logs/train_recall.csv` | Train R@K, AP, loss (subset, clean transforms) |
| **Plots** | auto | `{out}/logs/*.png` | train/val loss, R@K, AP, gates, temperature |
## Decision rule
| Delta R@1 (drone→satellite) | Verdict |
|---|---|
| >= +3% | **PASS** — captions informative, proceed to NADEZHDA teacher |
| +1% to +3% | MARGINAL — add VLM refinement, re-run |
| 0 to +1% | WEAK — redesign caption pipeline |
| < 0 | HARMFUL — critical bug |
## Code style
- `from __future__ import annotations` everywhere
- Type hints on all signatures
- Google-style docstrings
- English-only comments