From 7cc068e18ca5f90fc95b808eee26127115ceb948 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Tue, 21 Apr 2026 22:44:47 +0300 Subject: [PATCH] Update README: full architecture refresh (shared DINOv3, MONA bf16 top-12, 512-dim) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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) --- README.md | 105 ++++++++++++++++++++++++++++++------------------------ 1 file changed, 58 insertions(+), 47 deletions(-) diff --git a/README.md b/README.md index 5d4ab96..1db1faa 100644 --- a/README.md +++ b/README.md @@ -7,43 +7,45 @@ geo-localization (drone-to-satellite). ### Overview -Asymmetric dual-encoder with domain-specific frozen backbones and hierarchical -text fusion for drone-to-satellite image retrieval. +Shared DINOv3 WEB encoder with MONA adapters (last 12 blocks, bfloat16), +DGTRS-CLIP text fusion, and projection to 512-dim retrieval space. ``` ┌─────────────────────────────── QUERY BRANCH ─────────────────────────────────┐ │ │ -│ drone_img ──► DINOv3 ViT-L/16 LVD (frozen, 303M) ──► CLS token │ -│ [B,3,256,256] + MONA adapters (7M trainable) d_img [B,1024] │ -│ (2 per block × 24 blocks, bn=64) │ │ -│ │ │ +│ drone_img ──► DINOv3 ViT-L/16 WEB (frozen, 303M) ──► CLS [B,1024] │ +│ [B,3,256,256] + MONA adapters (3.5M, bf16) │ │ +│ (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] ─┐ │ │ L2 (full desc) ──► (frozen, 124M) ──► z₂ [B,768] ─┼─ cat ──[B,2304] │ │ L3 (fingerprint) ──► + LoRA r=4 (147K) ──► z₃ [B,768] ─┘ │ │ -│ (248 tok, KPS pos.) TextFusionMLP (shared, 3.4M) │ -│ Linear→GELU→Linear │ +│ (248 tok, KPS pos.) TextFusionMLP (shared, 1.5M) │ +│ + 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 / ‖q‖₂ ──► query [B,1024] │ +│ q̂ = q / ‖q‖₂ ──► query [B,512] │ └───────────────────────────────────────────────────────────────────────────────┘ ┌────────────────────────────── GALLERY BRANCH ────────────────────────────────┐ │ │ -│ sat_img ──► DINOv3 ViT-L/16 SAT (frozen, 303M) ──► CLS token │ -│ [B,3,256,256] + MONA adapters (7M trainable) s_img [B,1024] │ -│ (2 per block × 24 blocks, bn=64) │ │ -│ │ │ -│ sat_L1 ──► DGTRS-CLIP ──► z₁ [768] ─┐ │ -│ sat_L2 ──► (shared encoder)──► z₂ [768] ─┼─ cat ──► TextFusionMLP (shared) │ -│ sat_L3 ──► + LoRA ──► z₃ [768] ─┘ │ │ -│ s_txt [B,1024] │ -│ │ │ +│ sat_img ──► DINOv3 WEB (shared encoder) ──► CLS ──► Projection │ +│ [B,3,256,256] + MONA (shared adapters) s_img [B,512] │ +│ │ │ +│ sat_L1 ──► DGTRS-CLIP ──► z₁ ─┐ │ +│ sat_L2 ──► (shared) ──► z₂ ─┼─ cat ──► TextFusionMLP ──► s_txt [B,512] │ +│ sat_L3 ──► + LoRA ──► z₃ ─┘ │ │ +│ │ │ │ g = σ(α_g)·s_img + (1−σ(α_g))·s_txt GatedFusion_g │ │ │ │ -│ ĝ = g / ‖g‖₂ ──► gallery [B,1024] │ +│ ĝ = g / ‖g‖₂ ──► gallery [B,512] │ └───────────────────────────────────────────────────────────────────────────────┘ ┌────────────────────────────── RETRIEVAL ──────────────────────────────────────┐ @@ -51,8 +53,8 @@ text fusion for drone-to-satellite image retrieval. │ 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) │ │ │ -│ Retrieval space: 1024-dim (DINOv3 native, no projection layers) │ -│ TextFusionMLP shared between query and gallery branches │ +│ Retrieval space: 512-dim (DINOv3 1024 → projection → 512) │ +│ 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) │ │ 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 ```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) @@ -244,13 +246,14 @@ backward instead of storing them. Enabled by default (`gradient_checkpointing=Tr | **Max batch_size** (RTX 4090, shared encoder) | **8** | **24** | | 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 | |:---:|:---:|:---:| -| 16 | 14.7 GB | OK | -| 24 | 20.3 GB | OK (recommended) | -| 32 | >24 GB | OOM | +| 32 | 16.1 GB | OK | +| 40 | 19.1 GB | OK | +| **48** | **21.8 GB** | **OK (recommended)** | +| 56 | >24 GB | OOM | ### 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 | |---------|:---:|:---:|:---:|:---:| -| Default | 24 | 1 | 24 | 23 | -| Large effective batch | 24 | 4 | 96 | 23 per micro-batch | +| Default | 48 | 1 | 48 | 47 | +| Large effective batch | 48 | 4 | 192 | 47 per micro-batch | **Note:** gradient accumulation averages gradients across micro-batches, but each 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). ```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 \ - --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 @@ -288,18 +291,20 @@ Reported: R@1, R@5, R@10 for both q→g and g→q directions. ``` 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 - 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) + - Per optimizer step (accounts for gradient accumulation) - 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 +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 @@ -388,7 +393,8 @@ caption-test/ │ │ │ ├── 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) +│ │ │ ├── adapters.py # MONA (bf16) + LoRA adapters +│ │ ├── asymmetric_encoder.py # DINOv3 + projection + DGTRS + GatedFusion (v3) │ │ └── dual_encoder.py # GeoRSCLIP + GatedFusion (v2) │ ├── losses/ │ │ └── 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) ```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) +# With captions (L1/L2/L3, bs=48, 30 epochs, eval every epoch) 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) 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) ```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 +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 --batch-size 48 ``` ### 4. Enable diagnostics @@ -467,15 +478,15 @@ python -m src.training.train_gtauav --filter-meta meta/seg_filter.json ```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 + --filter-meta meta/seg_filter.json --batch-size 48 --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' + --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