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

Update README: full architecture refresh (shared DINOv3, MONA bf16 top-12, 512-dim)

Browse Source 
- Rewrite architecture diagram: shared DINOv3 WEB, projection 1024→512,
  MONA bf16 last 12 blocks, gradient checkpointing
- Update VRAM table: bs=48 at 21.8 GB (was bs=24)
- Update model summary: 432M total, 5.6M trainable, retrieval dim=512
- Update workflow: bs=48, 30 epochs, eval every epoch
- Document bf16 vs fp16 NaN issue (gamma=1e-6 underflow)
- Document MONA vs LoRA rationale for CVGL
- Update optimizer section: gradient accumulation, mixed precision details

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 22:44:47 +03:00
parent e55734425c
commit 7cc068e18c
1 changed files with 58 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
README.md
View File

@@ -7,43 +7,45 @@ geo-localization (drone-to-satellite).
### Overview ### Overview
Asymmetric dual-encoder with domain-specific frozen backbones and hierarchical Shared DINOv3 WEB encoder with MONA adapters (last 12 blocks, bfloat16),
text fusion for drone-to-satellite image retrieval. DGTRS-CLIP text fusion, and projection to 512-dim retrieval space.
``` ```
┌─────────────────────────────── QUERY BRANCH ─────────────────────────────────┐ ┌─────────────────────────────── QUERY BRANCH ─────────────────────────────────┐
│ │ │ │
│ drone_img ──► DINOv3 ViT-L/16 LVD (frozen, 303M) ──► CLS token │ │ drone_img ──► DINOv3 ViT-L/16 WEB (frozen, 303M) ──► CLS [B,1024] │
│ [B,3,256,256] + MONA adapters (7M trainable) d_img [B,1024] │ │ [B,3,256,256] + MONA adapters (3.5M, bf16) │ │
│ (2 per block × 24 blocks, bn=64) │ │ │ (2 per block × last 12 of 24, bn=64) │ │
│ │ │ │ + gradient checkpointing Projection │
│ Linear(1024→512) │
│ d_img [B,512] │
│ │ │
│ L1 (overview) ──► DGTRS-CLIP ViT-L-14 ──► z₁ [B,768] ─┐ │ │ L1 (overview) ──► DGTRS-CLIP ViT-L-14 ──► z₁ [B,768] ─┐ │
│ L2 (full desc) ──► (frozen, 124M) ──► z₂ [B,768] ─┼─ cat ──[B,2304] │ │ L2 (full desc) ──► (frozen, 124M) ──► z₂ [B,768] ─┼─ cat ──[B,2304] │
│ L3 (fingerprint) ──► + LoRA r=4 (147K) ──► z₃ [B,768] ─┘ │ │ │ L3 (fingerprint) ──► + LoRA r=4 (147K) ──► z₃ [B,768] ─┘ │ │
│ (248 tok, KPS pos.) TextFusionMLP (shared, 3.4M) │ │ (248 tok, KPS pos.) TextFusionMLP (shared, 1.5M) │
│ Linear→GELU→Linear │ │ + gradient checkpointing Linear(2304→512)→GELU→ │
│ Linear(512→512) │
│ │ │ │ │ │
│ d_txt [B,1024] │ │ d_txt [B,512] │
│ │ │ │ │ │
│ q = σ(α_q)·d_img + (1−σ(α_q))·d_txt GatedFusion_q │ │ q = σ(α_q)·d_img + (1−σ(α_q))·d_txt GatedFusion_q │
│ │ │ │ │ │
│ q̂ = q / ‖q‖₂ ──► query [B,1024] │ │ q̂ = q / ‖q‖₂ ──► query [B,512] │
└───────────────────────────────────────────────────────────────────────────────┘ └───────────────────────────────────────────────────────────────────────────────┘
┌────────────────────────────── GALLERY BRANCH ────────────────────────────────┐ ┌────────────────────────────── GALLERY BRANCH ────────────────────────────────┐
│ │ │ │
│ sat_img ──► DINOv3 ViT-L/16 SAT (frozen, 303M) ──► CLS token │ │ sat_img ──► DINOv3 WEB (shared encoder) ──► CLS ──► Projection │
│ [B,3,256,256] + MONA adapters (7M trainable) s_img [B,1024] │ │ [B,3,256,256] + MONA (shared adapters) s_img [B,512] │
│ (2 per block × 24 blocks, bn=64) │ │ │ │ │
│ │ │ │ sat_L1 ──► DGTRS-CLIP ──► z₁ ─┐ │
│ sat_L1 ──► DGTRS-CLIP ──► z₁ [768] ─┐ │ │ sat_L2 ──► (shared) ──► z₂ ─┼─ cat ──► TextFusionMLP ──► s_txt [B,512] │
│ sat_L2 ──► (shared encoder)──► z₂ [768] ─┼─ cat ──► TextFusionMLP (shared) │ │ sat_L3 ──► + LoRA ──► z₃ ─┘ │ │
│ sat_L3 ──► + LoRA ──► z₃ [768] ─┘ │ │ │ │ │
│ s_txt [B,1024] │
│ │ │
│ g = σ(α_g)·s_img + (1−σ(α_g))·s_txt GatedFusion_g │ │ g = σ(α_g)·s_img + (1−σ(α_g))·s_txt GatedFusion_g │
│ │ │ │ │ │
│ ĝ = g / ‖g‖₂ ──► gallery [B,1024] │ │ ĝ = g / ‖g‖₂ ──► gallery [B,512] │
└───────────────────────────────────────────────────────────────────────────────┘ └───────────────────────────────────────────────────────────────────────────────┘
┌────────────────────────────── RETRIEVAL ──────────────────────────────────────┐ ┌────────────────────────────── RETRIEVAL ──────────────────────────────────────┐
@@ -51,8 +53,8 @@ text fusion for drone-to-satellite image retrieval.
│ similarity = q̂ · ĝᵀ / τ (τ learnable, init=0.07, clamp [0.01, 0.5]) │ │ 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) │ │ loss = 0.6·CE(q→g) + 0.4·CE(g→q) (label_smoothing=0.1) │
│ │ │ │
│ Retrieval space: 1024-dim (DINOv3 native, no projection layers) │ │ Retrieval space: 512-dim (DINOv3 1024 → projection → 512) │
│ TextFusionMLP shared between query and gallery branches │ │ All shared: one DINOv3, one MONA set, one DGTRS-CLIP, one TextFusionMLP │
│ For sat images without captions: s_txt=None → g = s_img (gate passthrough) │ │ For sat images without captions: s_txt=None → g = s_img (gate passthrough) │
│ BASELINE: σ(α) = 1.0 for both branches (text disabled, DGTRS not loaded) │ │ BASELINE: σ(α) = 1.0 for both branches (text disabled, DGTRS not loaded) │
└───────────────────────────────────────────────────────────────────────────────┘ └───────────────────────────────────────────────────────────────────────────────┘
@@ -91,7 +93,7 @@ adapted with **LoRA** (rank=4 on Q/V in all 12 blocks).
where `[z₁ ; z₂ ; z₃] ∈ ℝ^{B×2304}` is the concatenation of three 768-dim DGTRS-CLIP embeddings, and where `[z₁ ; z₂ ; z₃] ∈ ℝ^{B×2304}` is the concatenation of three 768-dim DGTRS-CLIP embeddings, and
```math ```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} \text{MLP}: \text{Linear}(2304, 512) \to \text{GELU} \to \text{Linear}(512, 512), \quad \mathbf{z}_{\text{text}} \in \mathbb{R}^{B \times 512}
``` ```
### Gated fusion (separate gates for query and gallery) ### Gated fusion (separate gates for query and gallery)
@@ -244,13 +246,14 @@ backward instead of storing them. Enabled by default (`gradient_checkpointing=Tr
| **Max batch_size** (RTX 4090, shared encoder) | **8** | **24** | | **Max batch_size** (RTX 4090, shared encoder) | **8** | **24** |
| Speed penalty | — | ~20-30% slower per step | | Speed penalty | — | ~20-30% slower per step |
VRAM tested on RTX 4090 (24 GB) with shared DINOv3 WEB + DGTRS-CLIP + text: VRAM tested on RTX 4090 (24 GB) with shared DINOv3 WEB + MONA top-12 bf16 + DGTRS-CLIP + text:
| `batch_size` | Peak VRAM | Status | | `batch_size` | Peak VRAM | Status |
|:---:|:---:|:---:| |:---:|:---:|:---:|
| 16 | 14.7 GB | OK | | 32 | 16.1 GB | OK |
| 24 | 20.3 GB | OK (recommended) | | 40 | 19.1 GB | OK |
| 32 | >24 GB | OOM | | **48** | **21.8 GB** | **OK (recommended)** |
| 56 | >24 GB | OOM |
### Gradient accumulation ### Gradient accumulation
@@ -262,17 +265,17 @@ effective_batch_size = batch_size × grad_accum_steps
| Setting | `batch_size` | `grad_accum_steps` | Effective batch | In-batch negatives | | Setting | `batch_size` | `grad_accum_steps` | Effective batch | In-batch negatives |
|---------|:---:|:---:|:---:|:---:| |---------|:---:|:---:|:---:|:---:|
| Default | 24 | 1 | 24 | 23 | | Default | 48 | 1 | 48 | 47 |
| Large effective batch | 24 | 4 | 96 | 23 per micro-batch | | Large effective batch | 48 | 4 | 192 | 47 per micro-batch |
**Note:** gradient accumulation averages gradients across micro-batches, but each **Note:** gradient accumulation averages gradients across micro-batches, but each
micro-batch still only sees `batch_size` in-batch negatives. To increase the number 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). of negatives per forward pass, increase `batch_size` directly (requires more VRAM).
```bash ```bash
# Example: effective batch of 96 with gradient accumulation # Example: effective batch of 192 with gradient accumulation
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \ python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
--filter-meta meta/seg_filter.json --batch-size 24 --grad-accum 4 --filter-meta meta/seg_filter.json --batch-size 48 --grad-accum 4
``` ```
### Metrics ### Metrics
@@ -288,18 +291,20 @@ Reported: R@1, R@5, R@10 for both q→g and g→q directions.
``` ```
Optimizer: AdamW Optimizer: AdamW
- MONA adapters, TextFusionMLP, gate α_q, gate α_g, logit_scale: - MONA adapters, projection, TextFusionMLP, gate α_q, gate α_g, logit_scale:
lr = 1e-4, weight_decay = 1e-4 lr = 1e-4, weight_decay = 1e-4
- LoRA adapters (DGTRS-CLIP text encoder): - LoRA adapters (DGTRS-CLIP text encoder):
lr = 1e-5 (10× lower, --text-lr-factor 0.1) lr = 1e-5 (10× lower, --text-lr-factor 0.1)
Scheduler: Linear warmup (2 epochs) + cosine annealing Scheduler: Linear warmup (2 epochs) + cosine annealing
- Per-step (not per-epoch) - Per optimizer step (accounts for gradient accumulation)
- warmup: lr linearly ramps from 0 to lr_max over warmup_steps - warmup: lr linearly ramps from 0 to lr_max over warmup_steps
- cosine: lr decays from lr_max to 0 over remaining steps - cosine: lr decays from lr_max to 0 over remaining steps
Gradient clipping: max_norm = 1.0 Gradient clipping: max_norm = 1.0
Mixed precision: AMP fp16 for model forward, fp32 for loss Gradient accumulation: configurable (default 1, --grad-accum N)
Mixed precision: AMP fp16 for DINOv3/DGTRS forward, bf16 for MONA, fp32 for loss
Gradient checkpointing: DINOv3 (24 blocks) + DGTRS-CLIP (12 blocks)
``` ```
### Augmentations ### Augmentations
@@ -388,7 +393,8 @@ caption-test/
│ │ │ ├── model.py # DGTRSTextEncoder, build_model, tokenize │ │ │ ├── model.py # DGTRSTextEncoder, build_model, tokenize
│ │ │ ├── simple_tokenizer.py # BPE tokenizer (248 tokens) │ │ │ ├── simple_tokenizer.py # BPE tokenizer (248 tokens)
│ │ │ └── bpe_simple_vocab_16e6.txt.gz │ │ │ └── bpe_simple_vocab_16e6.txt.gz
│ │ ├── asymmetric_encoder.py # DINOv3 + DGTRS + GatedFusion (v3) │ │ │ ├── adapters.py # MONA (bf16) + LoRA adapters
│ │ ├── asymmetric_encoder.py # DINOv3 + projection + DGTRS + GatedFusion (v3)
│ │ └── dual_encoder.py # GeoRSCLIP + GatedFusion (v2) │ │ └── dual_encoder.py # GeoRSCLIP + GatedFusion (v2)
│ ├── losses/ │ ├── losses/
│ │ └── multi_infonce.py # InfoNCE with learnable temperature │ │ └── multi_infonce.py # InfoNCE with learnable temperature
@@ -442,24 +448,29 @@ python -m scripts.filter_segmentation --output meta/seg_filter.json
### 2. Train with gin configs (recommended) ### 2. Train with gin configs (recommended)
```bash ```bash
# Baseline (no text, 10 epochs) # With captions (L1/L2/L3, bs=48, 30 epochs, eval every epoch)
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 \ python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
--filter-meta meta/seg_filter.json --filter-meta meta/seg_filter.json --batch-size 48 --epochs 30 \
--gin-param 'TrainConfigGTAUAV.eval_every=1'
# Baseline (no text)
python -m src.training.train_gtauav --config conf/gtauav_baseline.gin \
--filter-meta meta/seg_filter.json --batch-size 48 --epochs 30
# Text-heavy (gate=0.3, 70% text weight) # Text-heavy (gate=0.3, 70% text weight)
python -m src.training.train_gtauav --config conf/gtauav_text_heavy.gin \ python -m src.training.train_gtauav --config conf/gtauav_text_heavy.gin \
--filter-meta meta/seg_filter.json --filter-meta meta/seg_filter.json --batch-size 48 --epochs 30
# Image-heavy (gate=0.9, 10% text weight)
python -m src.training.train_gtauav --config conf/gtauav_image_heavy.gin \
--filter-meta meta/seg_filter.json --batch-size 48 --epochs 30
``` ```
### 3. Train without gin (CLI-only) ### 3. Train without gin (CLI-only)
```bash ```bash
python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.json python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.json --batch-size 48
python -m src.training.train_gtauav --filter-meta meta/seg_filter.json python -m src.training.train_gtauav --filter-meta meta/seg_filter.json --batch-size 48
``` ```
### 4. Enable diagnostics ### 4. Enable diagnostics
@@ -467,15 +478,15 @@ python -m src.training.train_gtauav --filter-meta meta/seg_filter.json
```bash ```bash
# W&B + Grad-CAM + PyTorch Profiler # W&B + Grad-CAM + PyTorch Profiler
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \ python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
--filter-meta meta/seg_filter.json --wandb --gradcam --profile --filter-meta meta/seg_filter.json --batch-size 48 --wandb --gradcam --profile
# Gin parameter overrides from CLI # Gin parameter overrides from CLI
python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \ python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
--filter-meta meta/seg_filter.json \ --filter-meta meta/seg_filter.json \
--gin-param 'TrainConfigGTAUAV.batch_size=16' 'TrainConfigGTAUAV.epochs=20' --gin-param 'TrainConfigGTAUAV.eval_every=1' 'TrainConfigGTAUAV.epochs=30'
``` ```
CLI flags (`--wandb`, `--gradcam`, `--profile`, `--epochs`, etc.) take priority over gin config. CLI flags (`--wandb`, `--gradcam`, `--profile`, `--epochs`, `--batch-size`, etc.) take priority over gin config.
### 5. Resume from checkpoint ### 5. Resume from checkpoint