From 11fc348ecde29054bedc411ada778ada033ef1db Mon Sep 17 00:00:00 2001 From: pikaliov Date: Tue, 21 Apr 2026 21:44:10 +0300 Subject: [PATCH] =?UTF-8?q?Add=20projection=201024=E2=86=92512,=20revert?= =?UTF-8?q?=20MONA=20bottleneck=20to=2064?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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) --- src/models/adapters.py | 4 ++-- src/models/asymmetric_encoder.py | 34 ++++++++++++++++++++------------ 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/src/models/adapters.py b/src/models/adapters.py index 9f5c576..d6bb944 100644 --- a/src/models/adapters.py +++ b/src/models/adapters.py @@ -55,7 +55,7 @@ class MonaAdapter(nn.Module): dropout: Dropout rate. """ - def __init__(self, dim: int = 1024, bottleneck: int = 32, dropout: float = 0.1) -> None: + def __init__(self, dim: int = 1024, bottleneck: int = 64, dropout: float = 0.1) -> None: super().__init__() self.down = nn.Linear(dim, bottleneck) self.up = nn.Linear(bottleneck, dim) @@ -146,7 +146,7 @@ class LoRALinear(nn.Module): def inject_mona_into_dinov3( model: nn.Module, - bottleneck: int = 32, + bottleneck: int = 64, dropout: float = 0.1, ) -> int: """Inject MONA adapters into a frozen DINOv3ViT model. diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 8a4aea3..07b633e 100644 --- a/src/models/asymmetric_encoder.py +++ b/src/models/asymmetric_encoder.py @@ -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)