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

248
README.md
View File

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