pikaliov 9200772bea Fix NaN loss: revert MONA to fp32, fix loss logging
- 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>
2026-04-21 22:09:49 +03:00

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

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

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

MONA runs in fp16 (AMP-native) with gradient checkpointing to save VRAM.

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

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

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

\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, ..., 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 + DGTRS-CLIP + text:

batch_size Peak VRAM Status
16 14.7 GB OK
24 20.3 GB OK (recommended)
32 >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 24 1 24 23
Large effective batch 24 4 96 23 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 96 with gradient accumulation
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
    --filter-meta meta/seg_filter.json --batch-size 24 --grad-accum 4

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.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, fp16) 6.85M 6.85M 2 per block × 24 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) 436M 9.0M (2.06%) retrieval dim = 512

Asymmetric mode (--shared-encoder false): uses separate DINOv3 WEB (drone) + DINOv3 SAT (satellite) encoders with independent MONA adapters. 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
# 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
Description
Caption quality test on UAV-VisLoc
Readme 10 MiB
Languages
Python 100%