From cb477f4b4089cea5595a8fc22074ac814cf789ad Mon Sep 17 00:00:00 2001 From: pikaliov Date: Fri, 24 Apr 2026 16:26:17 +0300 Subject: [PATCH] Simplify model: shared DINOv3 WEB + MONA in last 12/24 blocks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three related architecture changes, driven by a cost/simplicity trade-off: 1. **Shared encoder**: one DINOv3 LVD-1689M (WEB) processes both drone and satellite images. Previously asymmetric — separate WEB (drone) and SAT-493M (satellite) encoders. Saves ~303M frozen params and halves VRAM for the image tower. Expected to lose some satellite-domain inductive bias; MONA adapters pick up the slack. 2. **MONA in last 12/24 blocks**: adapters injected only in the top half of the ViT. The lowest 12 blocks keep their pretrained features untouched. Trainable MONA count drops from 14.0M (48 adapters × 2 encoders) to 3.5M (24 adapters × 1 encoder). 3. **No DINO_SAT**: `nn_models/DINO_SAT` is no longer loaded by the default config. It stays on disk and the path param is kept for backward compat with asymmetric checkpoints. Parameter counts (with text fusion + LoRA + gates): Before: 17.6M trainable / 733M total (2.35%) After: 7.06M trainable / 434M total (1.63%) Also fixes a pre-existing resume bug: checkpoints now record `shared_encoder`, `baseline_mode`, `mona_bottleneck`, `mona_last_n_blocks` so `AsymmetricEncoder.load_checkpoint` can rebuild the right architecture. Old checkpoints still load (missing keys fall back to asymmetric defaults via `ckpt.get(..., )`). Co-Authored-By: Claude Opus 4.7 (1M context) --- CLAUDE.md | 34 +++++++++++++++----------------- conf/gtauav_balanced.gin | 4 +++- src/models/asymmetric_encoder.py | 2 ++ src/training/train_gtauav.py | 13 ++++++++++-- 4 files changed, 32 insertions(+), 21 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index f244857..36d64bd 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,10 +1,13 @@ # Caption Quality Test for Cross-View Geo-Localization -## Архитектура системы (v3, 2026-04-21) — GTA-UAV эксперимент +## Архитектура системы (v3, 2026-04-24) — GTA-UAV эксперимент ``` +Shared DINOv3 ViT-L/16 (LVD-1689M, frozen + MONA in last 12/24 blocks) +для обеих веток — drone и satellite кодируются одним encoder. + QUERY BRANCH (drone + L1/L2/L3 captions): - drone_img [B,3,256,256] --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> d_img [B,1024] + drone_img [B,3,256,256] --> DINOv3 ViT-L/16 (shared) --> d_img [B,1024] | L1 --> DGTRS-CLIP (248 tok) --> z₁ [768] --\ | L2 --> DGTRS-CLIP (248 tok) --> z₂ [768] ---+-- cat --> MLP(2304→1024→1024) --> d_txt [B,1024] @@ -15,7 +18,7 @@ QUERY BRANCH (drone + L1/L2/L3 captions): q̂ = q/‖q‖₂ --> query [B,1024] GALLERY BRANCH (satellite + satellite captions): - sat_img [B,3,256,256] --> DINOv3 ViT-L/16 SAT-493M (frozen) --> s_img [B,1024] + sat_img [B,3,256,256] --> DINOv3 ViT-L/16 (shared, same weights) --> s_img [B,1024] | sat_L1 --> DGTRS-CLIP --> z₁ --\ | sat_L2 --> DGTRS-CLIP --> z₂ ---+-- cat --> MLP (shared) --> s_txt [B,1024] @@ -54,15 +57,16 @@ BASELINE: σ(α_q)=σ(α_g)=1.0, text disabled, DGTRS not loaded - Transformer: sequence-first (LND), nn.MultiheadAttention, 12 layers - Tokenizer: BPE SimpleTokenizer (248 tokens, vocab 49408) -### Trainable parameters: 17.6M из 748M (2.35%) -- **MONA adapters** (2×DINOv3): 14.0M (2 per block × 24 × 2 encoders, bottleneck=64) +### Trainable parameters: 7.06M из 434M (1.63%) +- **MONA adapters** (shared DINOv3): 3.5M (2 per block × 12 last blocks, bottleneck=64) - **LoRA** (DGTRS-CLIP): 147K (Q+V, rank=4, 12 blocks) - TextFusionMLP (shared): Linear(2304,1024)+GELU+Linear(1024,1024) = ~3.4M - gate α_q + α_g: 2 scalars - logit_scale: 1 scalar (learnable temperature) -- DINOv3 x2 + DGTRS: frozen backbone weights +- DINOv3 (1 encoder) + DGTRS: frozen backbone weights - **Без projection layers** — retrieval space = DINOv3 native 1024-dim - **AMP:** frozen layers fp16, adapters + loss fp32 +- **Примечание:** ранее была asymmetric setup (2×DINOv3 WEB+SAT, MONA во всех 24 блоках) с 17.6M trainable / 733M total. Упростили до shared + last-12 MONA. ### Optimizer & Scheduler - **AdamW** с per-group LR: projections lr=1e-4, text encoder lr=1e-5 @@ -123,20 +127,14 @@ Eval: Resize(256) + CenterCrop(256) + ImageNet normalization. ## Backbones (v3) -### DINOv3 ViT-L/16 — Drone (web pretrained) +### DINOv3 ViT-L/16 — Shared (web pretrained) - **Checkpoint:** `nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth` - **Arch:** ViT-L/16, 24 layers, 16 heads, hidden=1024, MLP=4096, 303M params - **Input:** 256x256, ImageNet normalization, patch=16 → 256 patches - **Register tokens:** 4, RoPE theta=100.0 -- **Status:** frozen - -### DINOv3 ViT-L/16 — Satellite (sat pretrained) -- **Checkpoint:** `nn_models/DINO_SAT/model.safetensors` -- **HuggingFace:** `facebook/dinov3-vitl16-pretrain-sat493m` -- **Arch:** идентична DINO_WEB (ViT-L/16, hidden=1024, 303M params) -- **Input:** 256x256 -- **Config:** `nn_models/DINO_SAT/config.json` — BROKEN (auth error), используем конфиг от DINO_WEB -- **Status:** frozen +- **MONA:** 24 адаптера в последних 12 блоках (blocks 12-23), bottleneck=64, 3.5M trainable +- **Status:** frozen кроме MONA +- **Примечание:** ранее asymmetric — использовался отдельно `nn_models/DINO_SAT/model.safetensors` (sat493m pretrain) для satellite ветки. Упростили до shared WEB-энкодера. ### DGTRS-CLIP ViT-L-14 (LRSCLIP) — Text encoder - **Checkpoint:** `nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt` @@ -306,10 +304,10 @@ python -m scripts.compare_runs \ ### V3 (GTA-UAV, DINOv3 ViT-L/16, 256x256) | Фаза | Оценка | |------|--------| -| VRAM: 2x DINOv3-L + LRSCLIP + batch 64 | ~18-22 GB | +| VRAM: DINOv3-L (shared) + LRSCLIP + batch 64 | ~10-14 GB (было ~18-22 с 2× DINOv3) | | GPU mem (smoke test, batch 4) | 3.1 GB | | Batch size | 64 (default) | -| Total params | 733M (10.9M trainable, 1.49%) | +| Total params | 434M (7.06M trainable, 1.63%) — shared encoder + MONA в last 12/24 blocks | ### V2 (UAV-GeoLoc, GeoRSCLIP) | Фаза | Время | diff --git a/conf/gtauav_balanced.gin b/conf/gtauav_balanced.gin index f379d81..f96518a 100644 --- a/conf/gtauav_balanced.gin +++ b/conf/gtauav_balanced.gin @@ -22,7 +22,9 @@ TrainConfigGTAUAV.device = "cuda" # ---- Model ---- TrainConfigGTAUAV.init_gate = 0.7 TrainConfigGTAUAV.baseline_mode = False -TrainConfigGTAUAV.shared_encoder = False +TrainConfigGTAUAV.shared_encoder = True # single DINOv3 WEB for both branches +TrainConfigGTAUAV.mona_bottleneck = 64 +TrainConfigGTAUAV.mona_last_n_blocks = 12 # inject MONA only in last 12/24 ViT blocks TrainConfigGTAUAV.gradient_checkpointing = True # ---- Loss ---- diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 23b63cd..275ed37 100644 --- a/src/models/asymmetric_encoder.py +++ b/src/models/asymmetric_encoder.py @@ -540,6 +540,8 @@ class AsymmetricEncoder(nn.Module): lrsclip_path=lrsclip_path, baseline_mode=ckpt.get("baseline_mode", False), shared_encoder=ckpt.get("shared_encoder", False), + mona_bottleneck=ckpt.get("mona_bottleneck", 64), + mona_last_n_blocks=ckpt.get("mona_last_n_blocks", 24), device=device, ) model.load_state_dict(ckpt["model_state"], strict=False) diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index f932c05..89048c2 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -86,7 +86,9 @@ class TrainConfigGTAUAV: lrsclip_path: str = _LRSCLIP init_gate: float = 0.7 baseline_mode: bool = False - shared_encoder: bool = False # asymmetric: WEB (drone) + SAT (satellite) + shared_encoder: bool = True # single DINOv3 WEB for both branches (simpler, half the params) + mona_bottleneck: int = 64 + mona_last_n_blocks: int = 12 # inject adapters only in last 12 of 24 ViT blocks gradient_checkpointing: bool = True # trade compute for VRAM (allows larger batch) # Training. @@ -537,6 +539,8 @@ def train(cfg: TrainConfigGTAUAV) -> None: init_gate=cfg.init_gate, baseline_mode=cfg.baseline_mode, shared_encoder=cfg.shared_encoder, + mona_bottleneck=cfg.mona_bottleneck, + mona_last_n_blocks=cfg.mona_last_n_blocks, device=cfg.device, ).to(cfg.device) LOGGER.info("embed_dim=%d", model.embed_dim) @@ -1006,13 +1010,18 @@ def train(cfg: TrainConfigGTAUAV) -> None: history.append(epoch_record) - # Save checkpoint. + # Save checkpoint. Model architecture flags go into the ckpt so + # `AsymmetricEncoder.load_checkpoint` can rebuild the right shape. _atomic_save( obj={ "epoch": epoch, "model_state": model.state_dict(), "optimizer_state": optimizer.state_dict(), "loss_state": loss_fn.state_dict(), + "baseline_mode": cfg.baseline_mode, + "shared_encoder": cfg.shared_encoder, + "mona_bottleneck": cfg.mona_bottleneck, + "mona_last_n_blocks": cfg.mona_last_n_blocks, }, path=output_dir / f"ckpt_epoch{epoch:03d}.pt", )