caption-test/README.md

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

Asymmetric dual-encoder with domain-specific frozen backbones and hierarchical text fusion for drone-to-satellite image retrieval.

┌─────────────────────────────── QUERY BRANCH ─────────────────────────────────┐
│                                                                               │
│  drone_img ──► DINOv3 ViT-L/16 LVD (frozen, 303M) ──► CLS token             │
│  [B,3,256,256]     + MONA adapters (7M trainable)       d_img [B,1024]       │
│                    (2 per block × 24 blocks, bn=64)         │                 │
│                                                             │                 │
│  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, 3.4M) │
│                                                 Linear→GELU→Linear           │
│                                                       │                      │
│                                                  d_txt [B,1024]              │
│                                                       │                      │
│              q = σ(α_q)·d_img + (1−σ(α_q))·d_txt    GatedFusion_q           │
│                            │                                                  │
│                      q̂ = q / ‖q‖₂  ──► query [B,1024]                        │
└───────────────────────────────────────────────────────────────────────────────┘

┌────────────────────────────── GALLERY BRANCH ────────────────────────────────┐
│                                                                               │
│  sat_img ──► DINOv3 ViT-L/16 SAT (frozen, 303M) ──► CLS token               │
│  [B,3,256,256]    + MONA adapters (7M trainable)      s_img [B,1024]         │
│                   (2 per block × 24 blocks, bn=64)        │                   │
│                                                           │                   │
│  sat_L1 ──► DGTRS-CLIP ──► z₁ [768] ─┐                                       │
│  sat_L2 ──► (shared encoder)──► z₂ [768] ─┼─ cat ──► TextFusionMLP (shared) │
│  sat_L3 ──► + LoRA         ──► z₃ [768] ─┘                │                  │
│                                                       s_txt [B,1024]         │
│                                                            │                  │
│         g = σ(α_g)·s_img + (1−σ(α_g))·s_txt    GatedFusion_g                │
│                            │                                                  │
│                     ĝ = g / ‖g‖₂  ──► gallery [B,1024]                       │
└───────────────────────────────────────────────────────────────────────────────┘

┌────────────────────────────── 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: 1024-dim (DINOv3 native, no projection layers)              │
│  TextFusionMLP shared between query and gallery branches                      │
│  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	15–30 tokens
L2	Full description	Complete P1 (inventory) + P2 (spatial layout with 5+ zones)	100–200 tokens
L3	Fingerprint	P3: unique landmarks and spatial signature for matching	20–50 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)

\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

\text{MLP}: \text{Linear}(2304, 1024) \to \text{GELU} \to \text{Linear}(1024, 1024), \quad \mathbf{z}_{\text{text}} \in \mathbb{R}^{B \times 1024}

Gated fusion (separate gates for query and gallery)

\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)}

\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×1024} — DINOv3+MONA image embeddings
• d_txt, s_txt ∈ ℝ^{B×1024} — fused text embeddings
• 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 (drone + sat)	After MSA and MLP in each of 24 blocks	7.0M per encoder
LoRA (rank=4)	DGTRS-CLIP text encoder	Q and V projections in all 12 blocks	147K

MONA adapter (per block):

\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)

\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}

LoRA (per attention layer):

\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

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):

\mathcal{L} = w_{q \to g} \cdot \mathcal{L}_{q \to g} + w_{g \to q} \cdot \mathcal{L}_{g \to q}

\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, ..., B−1] — positives on diagonal
• label smoothing = 0.1

Loss and adapters run in fp32 (AMP autocast disabled) to prevent gradient overflow.

Gradient accumulation

With batch_size=8 on a 24 GB GPU, VRAM is the bottleneck. 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	8	1	8	7
Recommended	8	8	64	7 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).

# Example: effective batch of 64 with 8 accumulation steps
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
    --filter-meta meta/seg_filter.json --grad-accum 8

Metrics

Metric	Formula	Direction
R@K (Recall at K)	fraction of queries where correct gallery is in top-K	drone → satellite (primary)
Delta R@1	R@1(with_text) − R@1(baseline)	higher = text helps

Reported: R@1, R@5, R@10 for both q→g and g→q directions.

Optimizer & scheduler

Optimizer: AdamW
  - MONA adapters, 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-step (not per-epoch)
  - 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
Mixed precision: AMP fp16 for model forward, fp32 for loss

Augmentations

Transform	Drone (train)	Satellite (train)	Eval
RandomResizedCrop(256, scale=0.7–1.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.1–2.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)	7.0M	7.0M	2 per block × 24 blocks, bottleneck=64
DGTRS-CLIP ViT-L-14 (text)	124M	frozen	backbone weights frozen
LoRA adapters (text)	147K	147K	Q+V, rank=4, 12 blocks
TextFusionMLP (shared)	3.4M	3.4M	Linear(2304,1024) + GELU + Linear(1024,1024)
GatedFusion α_q + α_g	2	2	separate gate scalars
logit_scale	1	1	learnable temperature
Total (shared)	438M	10.6M (2.42%)	retrieval dim = 1024

Asymmetric mode (--shared-encoder false): uses separate DINOv3 WEB (drone) + DINOv3 SAT (satellite) encoders with independent MONA adapters. Total: 748M params, 17.6M trainable. Requires ~4-5 GB more VRAM.

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 (10 epochs, v3)
│   ├── gtauav_baseline.gin       # GTA-UAV baseline, no text (v3)
│   ├── 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
│   │   ├── asymmetric_encoder.py # DINOv3 + 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

python -m scripts.make_split --output-dir meta
python -m scripts.filter_segmentation --output meta/seg_filter.json

2. Train with gin configs (recommended)

# Baseline (no text, 10 epochs)
python -m src.training.train_gtauav --config conf/gtauav_baseline.gin \
    --filter-meta meta/seg_filter.json

# With captions (L1/L2/L3, 10 epochs)
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
    --filter-meta meta/seg_filter.json

# 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

3. Train without gin (CLI-only)

python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.json
python -m src.training.train_gtauav --filter-meta meta/seg_filter.json

4. Enable diagnostics

# W&B + Grad-CAM + PyTorch Profiler
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
    --filter-meta meta/seg_filter.json --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.batch_size=16' 'TrainConfigGTAUAV.epochs=20'

CLI flags (--wandb, --gradcam, --profile, --epochs, etc.) take priority over gin config.

5. Resume from checkpoint

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

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

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 averages + seaborn PNG plots

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