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

Simplify model: shared DINOv3 WEB + MONA in last 12/24 blocks

Browse Source 
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(..., <default>)`).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-24 16:26:17 +03:00
parent d98d853455
commit cb477f4b40
4 changed files with 32 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/training/train_gtauav.py
View File

@@ -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",
)