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

57
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).
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)
# 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)
self.drone_encoder.eval()
self.sat_encoder.eval()
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