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

Asymmetric encoders + MONA all 24 blocks + 1024-dim + hard negatives

Browse Source 
Architecture changes:
- Asymmetric DINOv3: WEB (drone) + SAT (satellite) with separate MONA
- MONA on all 24 blocks per encoder (was last 12)
- Remove projection, native 1024-dim retrieval space (was 512)
- Total: 748M params, 17.6M trainable (2.35%)

Hard negative memory bank:
- MoCo-style FIFO queue of 4096 detached gallery embeddings
- Each batch: B in-batch + Q queue negatives in InfoNCE
- Queue updated after each forward pass

Training config:
- batch_size=8, grad_accum=8, effective_batch=64
- eval_every=1 (eval + train recall every epoch)
- Max bs=24 with grad checkpointing on RTX 4090

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-24 08:47:33 +03:00
parent 75a4350d18
commit 04d5307221
5 changed files with 142 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
src/models/asymmetric_encoder.py
View File

@@ -297,14 +297,14 @@ 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,
embed_dim: int = 512,
shared_encoder: bool = False,
mona_bottleneck: int = 64,
mona_last_n_blocks: int = 24,
lora_rank: int = 4,
device: str = "cuda",
) -> None:
super().__init__()
self.embed_dim = embed_dim
self.embed_dim = self.DINO_DIM # native 1024, no projection
self.baseline_mode = baseline_mode
self.shared_encoder = shared_encoder
self.device = device
@@ -313,20 +313,17 @@ class AsymmetricEncoder(nn.Module):
if shared_encoder:
self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path)
self._freeze(self.image_encoder)
inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck)
inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
LOGGER.info("Shared encoder mode: single DINOv3 WEB for drone + satellite")
else:
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)
inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
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)
@@ -335,11 +332,11 @@ class AsymmetricEncoder(nn.Module):
else:
self.text_encoder = None
# Shared text fusion MLP: 3×768 -> embed_dim (512).
# Shared text fusion MLP: 3×768 -> 1024 (native DINOv3 dim).
if not baseline_mode:
self.text_fusion = TextFusionMLP(
text_dim=self.TEXT_DIM,
out_dim=embed_dim,
out_dim=self.DINO_DIM,
)
# Separate gated fusion for query and gallery branches.
@@ -353,20 +350,16 @@ class AsymmetricEncoder(nn.Module):
module.eval()
def encode_drone(self, images: torch.Tensor) -> torch.Tensor:
"""Encode drone images with MONA adapters + projection. Returns [B, embed_dim]."""
"""Encode drone images with MONA adapters. Returns [B, 1024]."""
if self.shared_encoder:
x = self.image_encoder(images)
else:
x = self.drone_encoder(images)
return self.image_projection(x)
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 + projection. Returns [B, embed_dim]."""
"""Encode satellite images with MONA adapters. Returns [B, 1024]."""
if self.shared_encoder:
x = self.image_encoder(images)
else:
x = self.sat_encoder(images)
return self.image_projection(x)
return self.image_encoder(images)
return self.sat_encoder(images)
def encode_text_levels(
self,
@@ -459,7 +452,6 @@ 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")
@@ -494,8 +486,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),
embed_dim=ckpt.get("embed_dim", 512),
shared_encoder=ckpt.get("shared_encoder", False),
device=device,
)
model.load_state_dict(ckpt["model_state"], strict=False)