Files
caption-test/README.md
pikaliov 29a09349e7 Add ML diagnostics tooling (W&B, TensorBoard, Grad-CAM, profiler) and gin configs
- 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>
2026-04-21 20:30:50 +03:00

388 lines
16 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
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 ──► CLS token │
│ [B,3,256,256] (frozen, 303M) d_img [B,1024] │
│ │ │
│ L1 (overview) ──► DGTRS-CLIP ──► z₁ [B,768] ─┐ │
│ L2 (full desc) ──► DGTRS-CLIP ──► z₂ [B,768] ─┼─ cat ──► [B,2304]│
│ L3 (fingerprint) ──► DGTRS-CLIP ──► z₃ [B,768] ─┘ │ │
│ (248 tokens, KPS pos. emb.) MLP(2304→1024→1024) │
│ │ │
│ 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 ──► CLS token │
│ [B,3,256,256] (frozen, 303M) s_img [B,1024] │
│ │ │
│ sat_L1 ──► DGTRS-CLIP ──► z₁ [768] ─┐ │
│ sat_L2 ──► DGTRS-CLIP ──► z₂ [768] ─┼─ cat ──► MLP ──► s_txt [1024]│
│ sat_L3 ──► DGTRS-CLIP ──► z₃ [768] ─┘ (shared MLP) │
│ │ │
│ g = σ(α_g)·s_img + (1σ(α_g))·s_txt GatedFusion_g │
│ │ │
│ ĝ = g / ‖g‖₂ ──► gallery [B,1024]│
└───────────────────────────────────────────────────────────────────────┘
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)
```
### 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, 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)
```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×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):
```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}
```
**LoRA** (per 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`
### 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.
### 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 LVD (drone) | 303M | frozen | backbone weights frozen |
| MONA adapters (drone) | 7.0M | 7.0M | 2 per block × 24 blocks, bottleneck=64 |
| DINOv3 ViT-L/16 SAT (satellite) | 303M | frozen | backbone weights frozen |
| MONA adapters (satellite) | 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** | **748M** | **17.6M (2.35%)** | retrieval dim = 1024 |
## 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)
│ ├── 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
```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
# 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)
```bash
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
```bash
# 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
```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 + plots** | auto | `{out}/logs/` | train.csv, val.csv, PNG plots every epoch |
## 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