diff --git a/README.md b/README.md index 82826e6..8855a39 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/conf/gtauav_balanced.gin b/conf/gtauav_balanced.gin index 7a68cec..7fead23 100644 --- a/conf/gtauav_balanced.gin +++ b/conf/gtauav_balanced.gin @@ -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 diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index dac6c0c..2bdba23 100644 --- a/src/models/asymmetric_encoder.py +++ b/src/models/asymmetric_encoder.py @@ -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 diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 2b2754d..2152600 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -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)