Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update docs: full architecture with tensor shapes, formulas, optimizer details

Browse Source 
README: architecture diagram with tensor dimensions, L1/L2/L3 text hierarchy
description, text fusion formula, InfoNCE loss formula with learnable
temperature, metrics table, optimizer/scheduler details with per-group LR,
augmentation table, model parameter summary.

CLAUDE.md: updated to DGTRS-CLIP (official architecture), loss formula,
optimizer/scheduler details, text encoder architecture notes.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 18:47:38 +03:00
parent bcb01bcb6d
commit 219bb779eb
2 changed files with 248 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
CLAUDE.md
View File

@@ -4,37 +4,64 @@
``` ```
QUERY BRANCH (drone + L1/L2/L3 captions): QUERY BRANCH (drone + L1/L2/L3 captions):
drone_img --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> CLS [B,1024] drone_img [B,3,256,256] --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> CLS [B,1024]
| |
proj_drone (1024->512) proj_drone: Linear(1024,512)
| |
L1 (overview) --> LRSCLIP (248 tok) --> z_L1 [768] --\ L1 (overview) --> DGTRS-CLIP (248 tok) --> z₁ [B,768] --\
L2 (full desc) --> LRSCLIP (248 tok) --> z_L2 [768] ---+-- concat [B,2304] L2 (full desc) --> DGTRS-CLIP (248 tok) --> z₂ [B,768] ---+-- cat --> [B,2304]
L3 (fingerprint) --> LRSCLIP (248 tok) --> z_L3 [768] --/ | L3 (fingerprint) --> DGTRS-CLIP (248 tok) --> z₃ [B,768] --/ |
MLP(2304->768->512) MLP(2304→768→512)
| |
GatedFusion(drone_512, text_512) --> L2-norm --> query [B,512] q = σ(α)·d_img + (1−σ(α))·d_txt GatedFusion
|
q̂ = q/‖q‖₂ --> query [B,512]
GALLERY BRANCH (satellite only): GALLERY BRANCH (satellite only):
sat_img --> DINOv3 ViT-L/16 SAT-493M (frozen) --> CLS [B,1024] sat_img [B,3,256,256] --> DINOv3 ViT-L/16 SAT-493M (frozen) --> CLS [B,1024]
| |
proj_sat (1024->512) --> L2-norm --> gallery [B,512] 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) LOSS: L = 0.6·CE(q̂·ĝᵀ/τ, targets) + 0.4·CE(ĝ·q̂ᵀ/τ, targets)
BASELINE: gate = 1.0 (text ignored) τ = 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%) ### Trainable parameters: 10.9M из 733M (1.49%)
- proj_drone: 1024x512 = ~524K - proj_drone: Linear(1024,512) = ~524K
- proj_sat: 1024x512 = ~524K - proj_sat: Linear(1024,512) = ~524K
- TextFusionMLP: 2304->768->512 = ~2.2M - TextFusionMLP: Linear(2304,768)+GELU+Linear(768,512) = ~2.2M
- gate alpha: 1 scalar - gate alpha: 1 scalar
- LRSCLIP partial unfreeze (last block + ln_final + text_projection): ~7.6M - logit_scale: 1 scalar (learnable temperature)
- Backbones DINOv3 x2 (303M each): frozen - 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 ### Image input: 256x256
DINOv3 ViT-L/16 с patch_size=16 → 16x16=256 patches на 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 эксперимент ### Предыдущая архитектура (v2) — UAV-GeoLoc эксперимент
@@ -54,12 +81,16 @@ Resize(256) + CenterCrop(256) + ImageNet normalization.
| `scripts/compare_runs.py` | Markdown/JSON сравнение baseline vs caption runs | | `scripts/compare_runs.py` | Markdown/JSON сравнение baseline vs caption runs |
| `scripts/generate_captions.py` | Offline caption generation (template/VLM/hybrid) | | `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/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) | | `scripts/filter_segmentation.py` | Scan segm masks, output meta JSON (exclude >=90% bg+water) |
## Backbones (v3) ## Backbones (v3)
@@ -81,10 +112,12 @@ Resize(256) + CenterCrop(256) + ImageNet normalization.
### DGTRS-CLIP ViT-L-14 (LRSCLIP) — Text encoder ### DGTRS-CLIP ViT-L-14 (LRSCLIP) — Text encoder
- **Checkpoint:** `nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt` - **Checkpoint:** `nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt`
- **Text dim:** 768, max tokens: 248 (KPS stretched from CLIP's 77) - **Код:** `src/models/dgtrs/` — официальная архитектура из github.com/MitsuiChen14/DGTRS
- **Содержит:** полную CLIP модель (visual + text), используем только text encoder - **Text dim:** 768, max tokens: 248 (KPS: mask1 pos 0-19 frozen + mask2 pos 20-247 trainable)
- **GitHub:** `github.com/MitsuiChen14/DGTRS` - **Transformer:** 12 layers, 12 heads, sequence-first (LND), QuickGELU
- **Status:** partial unfreeze (last block + text_projection) - **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) ### GeoRSCLIP ViT-B/32 (v2, legacy)
- **Checkpoint:** `checkpoints/RS5M_ViT-B-32.pt` - **Checkpoint:** `checkpoints/RS5M_ViT-B-32.pt`

248
README.md
View File

@@ -3,52 +3,160 @@
Validate whether generated text captions improve retrieval R@1 in cross-view Validate whether generated text captions improve retrieval R@1 in cross-view
geo-localization (drone-to-satellite). 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 ## Experiments
### V3 — GTA-UAV + DINOv3 + LRSCLIP (active) ### V3 — GTA-UAV + DINOv3 + DGTRS-CLIP (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)
**Dataset:** GTA-UAV-LR (33K drone + 14K satellite, GTA V synthetic) **Dataset:** GTA-UAV-LR (33K drone + 14K satellite, GTA V synthetic)
- RGB: `/home/servml/Документы/datasets/GTA-UAV-LR/` - RGB: `/home/servml/Документы/datasets/GTA-UAV-LR/`
@@ -76,23 +184,32 @@ caption-test/
│ ├── balanced.gin │ ├── balanced.gin
│ ├── baseline_no_text.gin │ ├── baseline_no_text.gin
│ └── text_heavy.gin │ └── text_heavy.gin
├── nn_models/ # Pre-trained checkpoints (v3) ├── nn_models/ # Pre-trained checkpoints (v3, gitignored)
│ ├── DINO_WEB/ # DINOv3 ViT-L/16 LVD-1689M │ ├── DINO_WEB/ # DINOv3 ViT-L/16 LVD-1689M (.pth)
│ ├── DINO_SAT/ # DINOv3 ViT-L/16 SAT-493M │ ├── DINO_SAT/ # DINOv3 ViT-L/16 SAT-493M (.safetensors)
│ └── LRSCLIP/ # DGTRS-CLIP ViT-L-14 │ └── 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/ ├── 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 │ ├── compare_runs.py # Delta R@1 comparison report
│ └── generate_captions.py # Offline caption generation │ └── generate_captions.py # Offline caption generation
├── src/ ├── src/
│ ├── datasets/ │ ├── 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) │ │ └── visloc_with_captions.py # UAV-GeoLoc loader (v2)
│ ├── models/ │ ├── 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) │ │ └── dual_encoder.py # GeoRSCLIP + GatedFusion (v2)
│ ├── losses/ │ ├── losses/
│ │ └── multi_infonce.py # InfoNCE with cosine temperature │ │ └── multi_infonce.py # InfoNCE with learnable temperature
│ ├── training/ │ ├── training/
│ │ ├── train_gtauav.py # Training loop GTA-UAV (v3) │ │ ├── train_gtauav.py # Training loop GTA-UAV (v3)
│ │ └── train.py # Training loop UAV-GeoLoc (v2) │ │ └── train.py # Training loop UAV-GeoLoc (v2)
@@ -105,9 +222,11 @@ caption-test/
``` ```
torch>=2.0 torch>=2.0
open_clip_torch
safetensors safetensors
timm coloredlogs
tqdm
ftfy
regex
gin-config gin-config
Pillow Pillow
numpy 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 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 ```bash
python -m scripts.compare_runs \ python -m scripts.compare_runs \
@@ -145,16 +271,16 @@ python -m scripts.compare_runs \
## Decision rule ## Decision rule
| Delta R@1 (drone->satellite) | Verdict | | Delta R@1 (drone→satellite) | Verdict |
|---|---| |---|---|
| >= +3% | PASS -- captions informative, proceed to NADEZHDA teacher | | >= +3% | **PASS** — captions informative, proceed to NADEZHDA teacher |
| +1% to +3% | MARGINAL -- add VLM refinement, re-run | | +1% to +3% | MARGINAL — add VLM refinement, re-run |
| 0 to +1% | WEAK -- redesign caption pipeline | | 0 to +1% | WEAK — redesign caption pipeline |
| < 0 | HARMFUL -- critical bug | | < 0 | HARMFUL — critical bug |
## Code style ## Code style
- `from __future__ import annotations` everywhere - `from __future__ import annotations` everywhere
- Type hints on all signatures - Type hints on all signatures
- Google-style docstrings - Google-style docstrings
- No emojis in code, English-only comments - English-only comments