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

Remove projections (1024 native), add satellite text, dual GatedFusion

Browse Source 
Architecture changes:
- Removed proj_drone/proj_sat (1024→512): retrieval space is now
  DINOv3 native 1024-dim, no information loss from projection
- TextFusionMLP: 2304→1024→1024 (was 2304→768→512), shared between branches
- Gallery branch now uses satellite captions (L1/L2/L3) via shared TextFusionMLP
- Two separate GatedFusion gates: α_q (query) and α_g (gallery)
- For sat images without captions (~57%): gate passes image features through

Dataset changes:
- GTAUAVDataset now loads satellite captions from caption index
- collate_gtauav_batch includes sat_caption_l1/l2/l3

Training loop:
- Passes satellite captions to model forward
- Logs both gate_q and gate_g values

11.1M trainable / 734M total (1.51%)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 19:01:30 +03:00
parent 219bb779eb
commit 0c41c1f017
6 changed files with 172 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
src/models/asymmetric_encoder.py
View File

@@ -208,20 +208,19 @@ class DINOv3ViT(nn.Module):
class TextFusionMLP(nn.Module):
"""Fuse L1/L2/L3 text embeddings via concat + MLP.
[B, 3*text_dim] -> [B, proj_dim]
[B, 3*text_dim] -> [B, out_dim]
"""
def __init__(
self,
text_dim: int = 768,
hidden_dim: int = 768,
proj_dim: int = 512,
out_dim: int = 1024,
) -> None:
super().__init__()
self.mlp = nn.Sequential(
nn.Linear(3 * text_dim, hidden_dim),
nn.Linear(3 * text_dim, out_dim),
nn.GELU(),
nn.Linear(hidden_dim, proj_dim),
nn.Linear(out_dim, out_dim),
)
def forward(
@@ -238,7 +237,7 @@ class TextFusionMLP(nn.Module):
z_l3: L3 fingerprint [B, text_dim].
Returns:
Fused text embedding [B, proj_dim].
Fused text embedding [B, out_dim].
"""
cat = torch.cat([z_l1, z_l2, z_l3], dim=-1)
return self.mlp(cat)
@@ -249,18 +248,24 @@ class TextFusionMLP(nn.Module):
# ---------------------------------------------------------------------------
class AsymmetricEncoder(nn.Module):
"""Asymmetric dual encoder for CVGL with text fusion.
"""Asymmetric dual encoder for CVGL with text fusion on both branches.
Query branch: DINOv3 LVD (drone) + LRSCLIP (L1/L2/L3) -> GatedFusion -> query
Gallery branch: DINOv3 SAT (satellite) -> gallery
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]
No projection layers — retrieval space is DINOv3 native 1024-dim.
Text fusion MLP is shared between branches (same caption format).
Two separate GatedFusion gates (drone/sat may weight text differently).
For satellite images without captions, GatedFusion passes image features through
(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).
lrsclip_path: Path to DGTRS-CLIP checkpoint (text encoder).
proj_dim: Shared projection dimension.
init_gate: Initial fusion gate (image weight).
baseline_mode: If True, gate = 1.0 (text ignored).
baseline_mode: If True, gate = 1.0 (text ignored), DGTRS not loaded.
device: Torch device string.
"""
@@ -272,13 +277,12 @@ class AsymmetricEncoder(nn.Module):
dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
proj_dim: int = 512,
init_gate: float = 0.7,
baseline_mode: bool = False,
device: str = "cuda",
) -> None:
super().__init__()
self.proj_dim = proj_dim
self.embed_dim = self.DINO_DIM
self.baseline_mode = baseline_mode
self.device = device
@@ -296,24 +300,16 @@ class AsymmetricEncoder(nn.Module):
else:
self.text_encoder = None
# Projection heads.
self.proj_drone = ProjectionHead(
in_dim=self.DINO_DIM, out_dim=proj_dim, use_mlp=False,
)
self.proj_sat = ProjectionHead(
in_dim=self.DINO_DIM, out_dim=proj_dim, use_mlp=False,
)
# Text fusion (L1/L2/L3 -> proj_dim).
# Shared text fusion MLP: 3×768 -> 1024 (same format for drone & sat captions).
if not baseline_mode:
self.text_fusion = TextFusionMLP(
text_dim=self.TEXT_DIM,
hidden_dim=self.TEXT_DIM,
proj_dim=proj_dim,
out_dim=self.DINO_DIM,
)
# Gated fusion.
self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
# Separate gated fusion for query and gallery branches.
self.fusion_query = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
self.fusion_gallery = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
@staticmethod
def _freeze(module: nn.Module) -> None:
@@ -354,8 +350,17 @@ class AsymmetricEncoder(nn.Module):
l1_texts: list[str],
l2_texts: list[str],
l3_texts: list[str],
) -> torch.Tensor:
"""Encode L1/L2/L3 captions and fuse. Returns [B, proj_dim]."""
) -> torch.Tensor | None:
"""Encode L1/L2/L3 captions and fuse. Returns [B, 1024] or None.
Returns None if all captions are empty (no text available).
For mixed batches (some have captions, some don't), encodes all
and lets GatedFusion handle per-sample gating.
"""
# Check if any caption is non-empty.
if all(t == "" for t in l1_texts):
return None
z_l1 = self._encode_single_text(l1_texts)
z_l2 = self._encode_single_text(l2_texts)
z_l3 = self._encode_single_text(l3_texts)
@@ -374,45 +379,49 @@ class AsymmetricEncoder(nn.Module):
caption_l1: list[str] | None = None,
caption_l2: list[str] | None = None,
caption_l3: list[str] | None = None,
sat_caption_l1: list[str] | None = None,
sat_caption_l2: list[str] | None = None,
sat_caption_l3: list[str] | None = None,
) -> dict[str, torch.Tensor]:
"""Forward pass.
Args:
drone_img: Drone images [B, 3, 256, 256].
sat_img: Satellite images [B, 3, 256, 256].
caption_l1: L1 overview captions.
caption_l2: L2 full description captions.
caption_l3: L3 fingerprint captions.
caption_l1/l2/l3: Drone L1/L2/L3 captions.
sat_caption_l1/l2/l3: Satellite L1/L2/L3 captions.
Returns:
Dict with 'query' [B, proj_dim], 'gallery' [B, proj_dim], 'gate'.
Dict with 'query' [B, 1024], 'gallery' [B, 1024],
'gate_q', 'gate_g'.
"""
# Gallery: satellite only.
sat_feat = self.encode_satellite(sat_img)
gallery = self.proj_sat(sat_feat)
# Query: drone + optional text.
# Image features (frozen DINOv3).
drone_feat = self.encode_drone(drone_img)
drone_proj = self.proj_drone(drone_feat)
sat_feat = self.encode_satellite(sat_img)
text_proj = None
has_text = (
caption_l1 is not None
and caption_l2 is not None
and caption_l3 is not None
and not self.baseline_mode
)
if has_text:
text_proj = self.encode_text_levels(caption_l1, caption_l2, caption_l3)
# Query branch: drone + drone text.
drone_text = None
if (caption_l1 is not None and caption_l2 is not None
and caption_l3 is not None and not self.baseline_mode):
drone_text = self.encode_text_levels(caption_l1, caption_l2, caption_l3)
query = self.fusion(drone_proj, text_proj)
# Re-normalize after fusion.
query = self.fusion_query(drone_feat, drone_text)
query = F.normalize(query, dim=-1)
# Gallery branch: satellite + satellite text.
sat_text = None
if (sat_caption_l1 is not None and sat_caption_l2 is not None
and sat_caption_l3 is not None and not self.baseline_mode):
sat_text = self.encode_text_levels(sat_caption_l1, sat_caption_l2, sat_caption_l3)
gallery = self.fusion_gallery(sat_feat, sat_text)
gallery = F.normalize(gallery, dim=-1)
return {
"query": query,
"gallery": gallery,
"gate": self.fusion.gate_value,
"gate_q": self.fusion_query.gate_value,
"gate_g": self.fusion_gallery.gate_value,
}
def trainable_parameters(self) -> list[nn.Parameter]:
@@ -425,7 +434,6 @@ class AsymmetricEncoder(nn.Module):
path.parent.mkdir(parents=True, exist_ok=True)
ckpt = {
"model_state": self.state_dict(),
"proj_dim": self.proj_dim,
"baseline_mode": self.baseline_mode,
**extra,
}
@@ -460,7 +468,6 @@ class AsymmetricEncoder(nn.Module):
dino_web_path=dino_web_path,
dino_sat_path=dino_sat_path,
lrsclip_path=lrsclip_path,
proj_dim=ckpt.get("proj_dim", 512),
baseline_mode=ckpt.get("baseline_mode", False),
device=device,
)