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

Switch to shared DINOv3 WEB encoder (saves ~4-5 GB VRAM)

Browse Source 
- Single DINOv3 WEB for both drone and satellite branches (shared_encoder=True default)
- One set of MONA adapters instead of two: 7M trainable vs 14M
- Total params: 438M (was 748M), trainable: 10.6M (was 17.6M)
- Asymmetric mode still available via shared_encoder=False
- Add gradient accumulation (grad_accum_steps, --grad-accum CLI flag)
- Update model summary in README

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 21:25:46 +03:00
parent 46b1208891
commit da2d2ea90e
4 changed files with 53 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
README.md
View File

@@ -254,16 +254,18 @@ Mixed precision: AMP fp16 for model forward, fp32 for loss
| Component | Params | Trainable | Notes |
|-----------|--------|-----------|-------|
| DINOv3 ViT-L/16 LVD (drone) | 303M | frozen | backbone weights frozen |
| MONA adapters (drone) | 7.0M | 7.0M | 2 per block × 24 blocks, bottleneck=64 |
| DINOv3 ViT-L/16 SAT (satellite) | 303M | frozen | backbone weights frozen |
| MONA adapters (satellite) | 7.0M | 7.0M | 2 per block × 24 blocks, bottleneck=64 |
| DINOv3 ViT-L/16 WEB (shared) | 303M | frozen | single encoder for drone + satellite |
| MONA adapters (shared) | 7.0M | 7.0M | 2 per block × 24 blocks, bottleneck=64 |
| DGTRS-CLIP ViT-L-14 (text) | 124M | frozen | backbone weights frozen |
| LoRA adapters (text) | 147K | 147K | Q+V, rank=4, 12 blocks |
| TextFusionMLP (shared) | 3.4M | 3.4M | Linear(2304,1024) + GELU + Linear(1024,1024) |
| GatedFusion α_q + α_g | 2 | 2 | separate gate scalars |
| logit_scale | 1 | 1 | learnable temperature |
| **Total** | **748M** | **17.6M (2.35%)** | retrieval dim = 1024 |
| **Total (shared)** | **438M** | **10.6M (2.42%)** | retrieval dim = 1024 |
> **Asymmetric mode** (`--shared-encoder false`): uses separate DINOv3 WEB (drone) + DINOv3 SAT
> (satellite) encoders with independent MONA adapters. Total: 748M params, 17.6M trainable.
> Requires ~4-5 GB more VRAM.
## Experiments

1
conf/gtauav_balanced.gin
View File

@@ -25,6 +25,7 @@ TrainConfigGTAUAV.device = "cuda"
# ---- Model ----
TrainConfigGTAUAV.init_gate = 0.7
TrainConfigGTAUAV.baseline_mode = False
TrainConfigGTAUAV.shared_encoder = True
# ---- Loss ----
TrainConfigGTAUAV.tau_init = 0.07

41
src/models/asymmetric_encoder.py
View File

@@ -249,10 +249,16 @@ class TextFusionMLP(nn.Module):
# ---------------------------------------------------------------------------
class AsymmetricEncoder(nn.Module):
"""Asymmetric dual encoder for CVGL with text fusion on both branches.
"""Dual encoder for CVGL with text fusion on both branches.
Query branch: DINOv3 LVD (drone) + text(L1/L2/L3) -> GatedFusion_q -> query [1024]
Gallery branch: DINOv3 SAT (sat) + text(L1/L2/L3) -> GatedFusion_g -> gallery [1024]
Supports two modes:
- **shared** (default): single DINOv3 WEB encoder for both drone and satellite,
one set of MONA adapters. Saves ~4-5 GB VRAM and halves adapter params.
- **asymmetric**: separate DINOv3 encoders (LVD for drone, SAT for satellite),
each with their own MONA adapters (legacy mode).
Query branch: DINOv3 (drone) + text(L1/L2/L3) -> GatedFusion_q -> query [1024]
Gallery branch: DINOv3 (sat) + text(L1/L2/L3) -> GatedFusion_g -> gallery [1024]
No projection layers — retrieval space is DINOv3 native 1024-dim.
Text fusion MLP is shared between branches (same caption format).
@@ -262,11 +268,12 @@ class AsymmetricEncoder(nn.Module):
(text_feat=None → gate acts as identity).
Args:
dino_web_path: Path to DINOv3 LVD checkpoint (drone encoder).
dino_sat_path: Path to DINOv3 SAT checkpoint (satellite encoder).
dino_web_path: Path to DINOv3 LVD checkpoint (used for both branches in shared mode).
dino_sat_path: Path to DINOv3 SAT checkpoint (only used in asymmetric mode).
lrsclip_path: Path to DGTRS-CLIP checkpoint (text encoder).
init_gate: Initial fusion gate (image weight).
baseline_mode: If True, gate = 1.0 (text ignored), DGTRS not loaded.
shared_encoder: If True, use single DINOv3 WEB for both branches.
device: Torch device string.
"""
@@ -280,6 +287,7 @@ class AsymmetricEncoder(nn.Module):
lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
init_gate: float = 0.7,
baseline_mode: bool = False,
shared_encoder: bool = True,
mona_bottleneck: int = 64,
lora_rank: int = 4,
device: str = "cuda",
@@ -287,15 +295,25 @@ class AsymmetricEncoder(nn.Module):
super().__init__()
self.embed_dim = self.DINO_DIM
self.baseline_mode = baseline_mode
self.shared_encoder = shared_encoder
self.device = device
# Image encoders (frozen + MONA adapters).
# Image encoder(s) (frozen + MONA adapters).
if shared_encoder:
# Single DINOv3 WEB for both branches — saves ~4-5 GB VRAM.
self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path)
self._freeze(self.image_encoder)
inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck)
LOGGER.info("Shared encoder mode: single DINOv3 WEB for drone + satellite")
else:
# Separate encoders (asymmetric mode).
self.drone_encoder = DINOv3ViT.from_pretrained(dino_web_path)
self.sat_encoder = DINOv3ViT.from_pretrained(dino_sat_path)
self._freeze(self.drone_encoder)
self._freeze(self.sat_encoder)
inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck)
inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck)
LOGGER.info("Asymmetric encoder mode: DINOv3 WEB (drone) + DINOv3 SAT (satellite)")
# Text encoder — official DGTRS architecture (frozen + LoRA).
if not baseline_mode:
@@ -324,10 +342,14 @@ class AsymmetricEncoder(nn.Module):
def encode_drone(self, images: torch.Tensor) -> torch.Tensor:
"""Encode drone images with MONA adapters. Returns [B, DINO_DIM]."""
if self.shared_encoder:
return self.image_encoder(images)
return self.drone_encoder(images)
def encode_satellite(self, images: torch.Tensor) -> torch.Tensor:
"""Encode satellite images with MONA adapters. Returns [B, DINO_DIM]."""
if self.shared_encoder:
return self.image_encoder(images)
return self.sat_encoder(images)
def encode_text_levels(
@@ -420,6 +442,7 @@ class AsymmetricEncoder(nn.Module):
ckpt = {
"model_state": self.state_dict(),
"baseline_mode": self.baseline_mode,
"shared_encoder": self.shared_encoder,
**extra,
}
tmp = path.with_suffix(path.suffix + ".tmp")
@@ -454,6 +477,7 @@ class AsymmetricEncoder(nn.Module):
dino_sat_path=dino_sat_path,
lrsclip_path=lrsclip_path,
baseline_mode=ckpt.get("baseline_mode", False),
shared_encoder=ckpt.get("shared_encoder", True),
device=device,
)
model.load_state_dict(ckpt["model_state"], strict=False)
@@ -464,11 +488,12 @@ class AsymmetricEncoder(nn.Module):
def train(self, mode: bool = True) -> AsymmetricEncoder:
"""Override to keep frozen encoders in eval mode."""
super().train(mode)
if self.shared_encoder:
self.image_encoder.eval()
else:
self.drone_encoder.eval()
self.sat_encoder.eval()
if self.text_encoder is not None:
# Text encoder partially unfrozen — set to train mode
# but frozen layers won't update anyway.
self.text_encoder.train(mode)
return self

5
src/training/train_gtauav.py
View File

@@ -73,6 +73,7 @@ class TrainConfigGTAUAV:
lrsclip_path: str = _LRSCLIP
init_gate: float = 0.7
baseline_mode: bool = False
shared_encoder: bool = True # single DINOv3 WEB for both branches (saves ~4-5 GB)
# Training.
resume_from: str | None = None # path to checkpoint for resuming
@@ -341,13 +342,15 @@ def train(cfg: TrainConfigGTAUAV) -> None:
start_epoch = resume_ckpt.get("epoch", -1) + 1
else:
mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)"
LOGGER.info("Building model — %s", mode_str)
enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)"
LOGGER.info("Building model — %s, %s", mode_str, enc_str)
model = AsymmetricEncoder(
dino_web_path=cfg.dino_web_path,
dino_sat_path=cfg.dino_sat_path,
lrsclip_path=cfg.lrsclip_path,
init_gate=cfg.init_gate,
baseline_mode=cfg.baseline_mode,
shared_encoder=cfg.shared_encoder,
device=cfg.device,
).to(cfg.device)