diff --git a/CLAUDE.md b/CLAUDE.md index 2d28599..4e9dbd3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -4,37 +4,64 @@ ``` QUERY BRANCH (drone + L1/L2/L3 captions): - drone_img --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> CLS [B,1024] - | - proj_drone (1024->512) - | - L1 (overview) --> LRSCLIP (248 tok) --> z_L1 [768] --\ - L2 (full desc) --> LRSCLIP (248 tok) --> z_L2 [768] ---+-- concat [B,2304] - L3 (fingerprint) --> LRSCLIP (248 tok) --> z_L3 [768] --/ | - MLP(2304->768->512) - | - GatedFusion(drone_512, text_512) --> L2-norm --> query [B,512] + drone_img [B,3,256,256] --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> CLS [B,1024] + | + proj_drone: Linear(1024,512) + | + L1 (overview) --> DGTRS-CLIP (248 tok) --> z₁ [B,768] --\ + L2 (full desc) --> DGTRS-CLIP (248 tok) --> z₂ [B,768] ---+-- cat --> [B,2304] + L3 (fingerprint) --> DGTRS-CLIP (248 tok) --> z₃ [B,768] --/ | + MLP(2304→768→512) + | + q = σ(α)·d_img + (1−σ(α))·d_txt GatedFusion + | + q̂ = q/‖q‖₂ --> query [B,512] GALLERY BRANCH (satellite only): - sat_img --> DINOv3 ViT-L/16 SAT-493M (frozen) --> CLS [B,1024] - | - proj_sat (1024->512) --> L2-norm --> gallery [B,512] + sat_img [B,3,256,256] --> DINOv3 ViT-L/16 SAT-493M (frozen) --> CLS [B,1024] + | + proj_sat: Linear(1024,512) + | + ĝ = g/‖g‖₂ --> gallery [B,512] -LOSS: InfoNCE(query, gallery) — symmetric, asymmetric weights (0.6 q->g, 0.4 g->q) -BASELINE: gate = 1.0 (text ignored) +LOSS: L = 0.6·CE(q̂·ĝᵀ/τ, targets) + 0.4·CE(ĝ·q̂ᵀ/τ, targets) + τ = 1/exp(logit_scale), learnable, clamped [0.01, 0.5], init=0.07 + label_smoothing=0.1 + +BASELINE: σ(α) = 1.0, text branch disabled, DGTRS not loaded ``` +### Text hierarchy (L1/L2/L3) +- **L1 overview:** первое предложение P1 — краткое описание land-cover (15-30 tok) +- **L2 full:** полные P1 + P2 — inventory + spatial layout (100-200 tok) +- **L3 fingerprint:** P3 — уникальные landmarks для matching (20-50 tok) +- **Fusion:** z_text = MLP([z₁; z₂; z₃]) — concat 3×768 → Linear(2304,768) → GELU → Linear(768,512) + +### Text encoder: DGTRS-CLIP (official architecture) +- Код: `src/models/dgtrs/` — из github.com/MitsuiChen14/DGTRS (Apache-2.0) +- KPS positional embedding: mask1 (pos 0-19, frozen) + mask2 (pos 20-247, trainable) +- Transformer: sequence-first (LND), nn.MultiheadAttention, 12 layers +- Tokenizer: BPE SimpleTokenizer (248 tokens, vocab 49408) + ### Trainable parameters: 10.9M из 733M (1.49%) -- proj_drone: 1024x512 = ~524K -- proj_sat: 1024x512 = ~524K -- TextFusionMLP: 2304->768->512 = ~2.2M +- proj_drone: Linear(1024,512) = ~524K +- proj_sat: Linear(1024,512) = ~524K +- TextFusionMLP: Linear(2304,768)+GELU+Linear(768,512) = ~2.2M - gate alpha: 1 scalar -- LRSCLIP partial unfreeze (last block + ln_final + text_projection): ~7.6M -- Backbones DINOv3 x2 (303M each): frozen +- logit_scale: 1 scalar (learnable temperature) +- DGTRS partial unfreeze (last resblock + ln_final + text_projection): ~7.6M +- DINOv3 x2 (303M each): frozen + +### Optimizer & Scheduler +- **AdamW** с per-group LR: projections lr=1e-4, text encoder lr=1e-5 +- **Linear warmup** (2 epochs) + **cosine annealing** (per-step) +- **Gradient clipping:** max_norm=1.0 +- **AMP:** fp16 для model forward, fp32 для loss (learnable temperature overflow fix) ### Image input: 256x256 DINOv3 ViT-L/16 с patch_size=16 → 16x16=256 patches на 256x256. -Resize(256) + CenterCrop(256) + ImageNet normalization. +Train: augmentations (drone: crop+flip+rot+jitter+blur, sat: crop+flip+jitter). +Eval: Resize(256) + CenterCrop(256) + ImageNet normalization. ### Предыдущая архитектура (v2) — UAV-GeoLoc эксперимент @@ -54,12 +81,16 @@ Resize(256) + CenterCrop(256) + ImageNet normalization. | `scripts/compare_runs.py` | Markdown/JSON сравнение baseline vs caption runs | | `scripts/generate_captions.py` | Offline caption generation (template/VLM/hybrid) | -### V3 (GTA-UAV, DINOv3 + LRSCLIP) — DONE +### V3 (GTA-UAV, DINOv3 + DGTRS-CLIP) — DONE | Файл | Назначение | |------|-----------| -| `src/models/asymmetric_encoder.py` | DINOv3ViT + LRSCLIPTextEncoder + TextFusionMLP + AsymmetricEncoder + GatedFusion | +| `src/models/dgtrs/model.py` | Официальная архитектура DGTRS-CLIP text encoder (Apache-2.0) | +| `src/models/dgtrs/simple_tokenizer.py` | BPE tokenizer (248 tokens, vocab 49408) | +| `src/models/asymmetric_encoder.py` | DINOv3ViT + TextFusionMLP + AsymmetricEncoder + GatedFusion | | `src/datasets/gtauav_dataset.py` | GTA-UAV-LR loader + L1/L2/L3 caption parsing из VLM JSON | -| `src/training/train_gtauav.py` | Training loop с eval, AMP, CLI args (--baseline, --filter-meta) | +| `src/losses/multi_infonce.py` | InfoNCE с learnable temperature (fp32), clamp [0.01, 0.5] | +| `src/training/train_gtauav.py` | Training loop с eval, AMP, per-group LR, warmup, --resume | +| `scripts/make_split.py` | 80/20 random split из всех пар | | `scripts/filter_segmentation.py` | Scan segm masks, output meta JSON (exclude >=90% bg+water) | ## Backbones (v3) @@ -81,10 +112,12 @@ Resize(256) + CenterCrop(256) + ImageNet normalization. ### DGTRS-CLIP ViT-L-14 (LRSCLIP) — Text encoder - **Checkpoint:** `nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt` -- **Text dim:** 768, max tokens: 248 (KPS stretched from CLIP's 77) -- **Содержит:** полную CLIP модель (visual + text), используем только text encoder -- **GitHub:** `github.com/MitsuiChen14/DGTRS` -- **Status:** partial unfreeze (last block + text_projection) +- **Код:** `src/models/dgtrs/` — официальная архитектура из github.com/MitsuiChen14/DGTRS +- **Text dim:** 768, max tokens: 248 (KPS: mask1 pos 0-19 frozen + mask2 pos 20-247 trainable) +- **Transformer:** 12 layers, 12 heads, sequence-first (LND), QuickGELU +- **Tokenizer:** BPE SimpleTokenizer (vocab 49408), 248 token context +- **Содержит:** полную CLIP модель (visual + text), используем только text encoder (124M params) +- **Status:** partial unfreeze (last resblock + ln_final + text_projection, ~7.6M trainable) ### GeoRSCLIP ViT-B/32 (v2, legacy) - **Checkpoint:** `checkpoints/RS5M_ViT-B-32.pt` diff --git a/README.md b/README.md index d1211a8..ad9bc0e 100644 --- a/README.md +++ b/README.md @@ -3,52 +3,160 @@ 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) [B,1024] │ +│ │ │ +│ proj_drone: Linear(1024,512) │ +│ │ │ +│ d_img [B,512] │ +│ │ │ +│ 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→768→512) │ +│ │ │ +│ d_txt [B,512] │ +│ │ │ +│ q = σ(α)·d_img + (1−σ(α))·d_txt GatedFusion │ +│ │ │ +│ q̂ = q / ‖q‖₂ ──► query [B,512] │ +└───────────────────────────────────────────────────────────────────────┘ + +┌──────────────────────────── GALLERY BRANCH ──────────────────────────┐ +│ │ +│ sat_img ──► DINOv3 ViT-L/16 SAT ──► CLS token │ +│ [B,3,256,256] (frozen, 303M) [B,1024] │ +│ │ │ +│ proj_sat: Linear(1024,512) │ +│ │ │ +│ ĝ = g / ‖g‖₂ ──► gallery [B,512] │ +└───────────────────────────────────────────────────────────────────────┘ + +BASELINE: σ(α) = 1.0 → q = d_img (text branch 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 | 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). + +**Text fusion:** + +``` +z_text = MLP( [z₁ ; z₂ ; z₃] ) + +where [z₁ ; z₂ ; z₃] ∈ ℝ^(B×2304) — concatenation of three 768-dim embeddings + MLP: Linear(2304, 768) → GELU → Linear(768, 512) + z_text ∈ ℝ^(B×512) +``` + +**Gated fusion:** + +``` +q = σ(α) · d_img + (1 − σ(α)) · d_txt + +where α — learnable scalar in logit-space (init: σ(α) ≈ 0.7) + σ — sigmoid function + d_img — projected drone image embedding [B, 512] + d_txt — fused text embedding [B, 512] +``` + +### Loss function + +Symmetric InfoNCE with learnable temperature (CLIP-style `logit_scale`): + +``` +L = w_q2g · L_q→g + w_g2q · L_g→q + +L_q→g = CrossEntropy( q̂ · ĝᵀ / τ, targets ) +L_g→q = CrossEntropy( ĝ · q̂ᵀ / τ, targets ) + +where τ = 1 / exp(logit_scale), logit_scale — learnable scalar + τ ∈ [0.01, 0.5] (clamped) + τ_init = 0.07 + w_q2g = 0.6, w_g2q = 0.4 + targets = [0, 1, 2, ..., B−1] (positives on diagonal) + label_smoothing = 0.1 +``` + +Loss is computed in **fp32** (outside AMP autocast) to prevent gradient overflow +in the learnable temperature. + +### 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 + - Projection heads (proj_drone, proj_sat, TextFusionMLP, gate α, logit_scale): + lr = 1e-4, weight_decay = 1e-4 + - DGTRS text encoder (last resblock + ln_final + text_projection): + 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 LVD (drone) | 303M | 0 | frozen | +| DINOv3 ViT-L/16 SAT (satellite) | 303M | 0 | frozen | +| DGTRS-CLIP ViT-L-14 (text) | 124M | ~7.6M | last block + ln_final + text_projection | +| proj_drone | 524K | 524K | Linear(1024, 512) | +| proj_sat | 524K | 524K | Linear(1024, 512) | +| TextFusionMLP | 2.2M | 2.2M | Linear(2304,768) + GELU + Linear(768,512) | +| GatedFusion α | 1 | 1 | scalar | +| logit_scale | 1 | 1 | learnable temperature | +| **Total** | **733M** | **10.9M (1.49%)** | | + ## Experiments -### V3 — GTA-UAV + DINOv3 + LRSCLIP (active) - -Asymmetric architecture with domain-specific image encoders and hierarchical text. - -``` -┌─────────────────────────── QUERY BRANCH ───────────────────────────┐ -│ │ -│ drone_img ──► DINOv3 ViT-L/16 LVD ──► CLS [1024] │ -│ (frozen, 303M) │ │ -│ proj_drone [1024→512] │ -│ │ │ -│ L1 (overview) ──► DGTRS-CLIP ──► [768] ─┐ │ -│ L2 (full desc) ──► DGTRS-CLIP ──► [768] ─┼─ concat [2304] │ -│ L3 (fingerprint) ──► DGTRS-CLIP ──► [768] ─┘ │ │ -│ (248 tok) MLP [2304→768→512] │ -│ │ │ -│ GatedFusion(img_512, text_512) │ -│ │ │ -│ L2-norm ──► query [512] │ -└─────────────────────────────────────────────────────────────────────┘ - -┌─────────────────────────── GALLERY BRANCH ─────────────────────────┐ -│ │ -│ sat_img ──► DINOv3 ViT-L/16 SAT ──► CLS [1024] │ -│ (frozen, 303M) │ │ -│ proj_sat [1024→512] │ -│ │ │ -│ L2-norm ──► gallery [512] │ -└─────────────────────────────────────────────────────────────────────┘ - -LOSS: InfoNCE(query, gallery) — symmetric, learnable temperature - weights: 0.6 × q→g + 0.4 × g→q - -BASELINE: gate = 1.0 (text branch disabled, no DGTRS loaded) -``` - -**Models:** -- Drone: DINOv3 ViT-L/16 (LVD-1689M, web pretrained) — 1024-dim, 303M params, frozen -- Satellite: DINOv3 ViT-L/16 (SAT-493M, satellite pretrained) — 1024-dim, 303M params, frozen -- Text: DGTRS-CLIP ViT-L-14 (LRSCLIP, 248 tokens) — 768-dim, partial unfreeze -- Total: 733M params, 10.9M trainable (1.49%) - -**Input:** 256x256, ImageNet normalization -**Training:** learnable temperature (CLIP logit_scale), per-group LR (proj 1e-4 / text 1e-5), warmup 2 epochs + cosine, augmentations (drone: crop+flip+rot+jitter+blur, sat: crop+flip+jitter) +### 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/` @@ -76,23 +184,32 @@ caption-test/ │ ├── balanced.gin │ ├── baseline_no_text.gin │ └── text_heavy.gin -├── nn_models/ # Pre-trained checkpoints (v3) -│ ├── DINO_WEB/ # DINOv3 ViT-L/16 LVD-1689M -│ ├── DINO_SAT/ # DINOv3 ViT-L/16 SAT-493M -│ └── LRSCLIP/ # DGTRS-CLIP ViT-L-14 +├── 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/ -│ ├── filter_segmentation.py # Meta-file: exclude 90%+ background/water +│ ├── 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 captions (v3) +│ │ ├── gtauav_dataset.py # GTA-UAV-LR loader + L1/L2/L3 parsing (v3) │ │ └── visloc_with_captions.py # UAV-GeoLoc loader (v2) │ ├── models/ -│ │ ├── asymmetric_encoder.py # DINOv3 + LRSCLIP + GatedFusion (v3) +│ │ ├── 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 cosine temperature +│ │ └── multi_infonce.py # InfoNCE with learnable temperature │ ├── training/ │ │ ├── train_gtauav.py # Training loop GTA-UAV (v3) │ │ └── train.py # Training loop UAV-GeoLoc (v2) @@ -105,9 +222,11 @@ caption-test/ ``` torch>=2.0 -open_clip_torch safetensors -timm +coloredlogs +tqdm +ftfy +regex gin-config Pillow numpy @@ -134,7 +253,14 @@ python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.jso python -m src.training.train_gtauav --filter-meta meta/seg_filter.json ``` -### 4. Compare and get verdict +### 4. Resume from checkpoint + +```bash +python -m src.training.train_gtauav --resume out/gtauav/with_text/ckpt_epoch004.pt \ + --filter-meta meta/seg_filter.json +``` + +### 5. Compare and get verdict ```bash python -m scripts.compare_runs \ @@ -145,16 +271,16 @@ python -m scripts.compare_runs \ ## Decision rule -| Delta R@1 (drone->satellite) | Verdict | +| 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 | +| >= +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 -- No emojis in code, English-only comments +- English-only comments