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

Add projection 1024→512, revert MONA bottleneck to 64

Browse Source 
- Add image_projection: Linear(1024→512) after DINOv3 CLS output
- Retrieval space: 512-dim (was 1024-dim native DINOv3)
- TextFusionMLP: 3×768→512→512 (was →1024→1024)
- GatedFusion operates in 512-dim
- MONA bottleneck restored to 64 (works at native 1024 inside DINOv3)
- Trainable: 8.9M (projection 0.5M + MONA 6.8M + LoRA 0.1K + MLP 1.5M)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 21:44:10 +03:00
parent d279a6e745
commit 11fc348ecd
2 changed files with 23 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
src/models/asymmetric_encoder.py
View File

@@ -298,25 +298,24 @@ class AsymmetricEncoder(nn.Module):
init_gate: float = 0.7,
baseline_mode: bool = False,
shared_encoder: bool = True,
mona_bottleneck: int = 32,
embed_dim: int = 512,
mona_bottleneck: int = 64,
lora_rank: int = 4,
device: str = "cuda",
) -> None:
super().__init__()
self.embed_dim = self.DINO_DIM
self.embed_dim = embed_dim
self.baseline_mode = baseline_mode
self.shared_encoder = shared_encoder
self.device = device
# 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)
@@ -325,6 +324,9 @@ class AsymmetricEncoder(nn.Module):
inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck)
LOGGER.info("Asymmetric encoder mode: DINOv3 WEB (drone) + DINOv3 SAT (satellite)")
# Projection: DINOv3 1024-dim -> embed_dim (512).
self.image_projection = nn.Linear(self.DINO_DIM, embed_dim)
# Text encoder — official DGTRS architecture (frozen + LoRA).
if not baseline_mode:
self.text_encoder = load_dgtrs_text_encoder(lrsclip_path)
@@ -333,11 +335,11 @@ class AsymmetricEncoder(nn.Module):
else:
self.text_encoder = None
# Shared text fusion MLP: 3×768 -> 1024 (same format for drone & sat captions).
# Shared text fusion MLP: 3×768 -> embed_dim (512).
if not baseline_mode:
self.text_fusion = TextFusionMLP(
text_dim=self.TEXT_DIM,
out_dim=self.DINO_DIM,
out_dim=embed_dim,
)
# Separate gated fusion for query and gallery branches.
@@ -351,16 +353,20 @@ class AsymmetricEncoder(nn.Module):
module.eval()
def encode_drone(self, images: torch.Tensor) -> torch.Tensor:
"""Encode drone images with MONA adapters. Returns [B, DINO_DIM]."""
"""Encode drone images with MONA adapters + projection. Returns [B, embed_dim]."""
if self.shared_encoder:
return self.image_encoder(images)
return self.drone_encoder(images)
x = self.image_encoder(images)
else:
x = self.drone_encoder(images)
return self.image_projection(x)
def encode_satellite(self, images: torch.Tensor) -> torch.Tensor:
"""Encode satellite images with MONA adapters. Returns [B, DINO_DIM]."""
"""Encode satellite images with MONA adapters + projection. Returns [B, embed_dim]."""
if self.shared_encoder:
return self.image_encoder(images)
return self.sat_encoder(images)
x = self.image_encoder(images)
else:
x = self.sat_encoder(images)
return self.image_projection(x)
def encode_text_levels(
self,
@@ -409,7 +415,7 @@ class AsymmetricEncoder(nn.Module):
sat_caption_l1/l2/l3: Satellite L1/L2/L3 captions.
Returns:
Dict with 'query' [B, 1024], 'gallery' [B, 1024],
Dict with 'query' [B, embed_dim], 'gallery' [B, embed_dim],
'gate_q', 'gate_g'.
"""
# Image features (frozen DINOv3).
@@ -453,6 +459,7 @@ class AsymmetricEncoder(nn.Module):
"model_state": self.state_dict(),
"baseline_mode": self.baseline_mode,
"shared_encoder": self.shared_encoder,
"embed_dim": self.embed_dim,
**extra,
}
tmp = path.with_suffix(path.suffix + ".tmp")
@@ -488,6 +495,7 @@ class AsymmetricEncoder(nn.Module):
lrsclip_path=lrsclip_path,
baseline_mode=ckpt.get("baseline_mode", False),
shared_encoder=ckpt.get("shared_encoder", True),
embed_dim=ckpt.get("embed_dim", 512),
device=device,
)
model.load_state_dict(ckpt["model_state"], strict=False)