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

25
src/datasets/gtauav_dataset.py
View File

@@ -176,13 +176,20 @@ class GTAUAVDataset(Dataset):
else:
continue # No match, skip.
# Get captions.
# Get drone captions.
cap_data = self.caption_index.get(drone_name)
if cap_data is not None:
l1, l2, l3 = _parse_caption_levels(cap_data["output"])
else:
l1 = l2 = l3 = _EMPTY_CAPTION
# Pre-parse satellite captions for all candidates.
sat_captions: dict[str, tuple[str, str, str]] = {}
for sat_name in sat_candidates:
sat_cap = self.caption_index.get(sat_name)
if sat_cap is not None:
sat_captions[sat_name] = _parse_caption_levels(sat_cap["output"])
self.entries.append({
"drone_name": drone_name,
"drone_dir": pair["drone_img_dir"],
@@ -192,6 +199,7 @@ class GTAUAVDataset(Dataset):
"caption_l1": l1,
"caption_l2": l2,
"caption_l3": l3,
"sat_captions": sat_captions,
})
def _load_image(self, directory: str, filename: str, transform: Callable | None = None) -> torch.Tensor:
@@ -222,18 +230,28 @@ class GTAUAVDataset(Dataset):
sat_img = self._load_image(entry["sat_dir"], sat_name, self.sat_transform)
# Captions with optional dropout.
# Drone captions with optional dropout.
if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
l1 = l2 = l3 = _EMPTY_CAPTION
else:
l1, l2, l3 = entry["caption_l1"], entry["caption_l2"], entry["caption_l3"]
# Satellite captions (empty string if not available).
sat_caps = entry["sat_captions"].get(sat_name)
if sat_caps is not None:
sat_l1, sat_l2, sat_l3 = sat_caps
else:
sat_l1 = sat_l2 = sat_l3 = _EMPTY_CAPTION
return {
"drone_img": drone_img,
"sat_img": sat_img,
"caption_l1": l1,
"caption_l2": l2,
"caption_l3": l3,
"sat_caption_l1": sat_l1,
"sat_caption_l2": sat_l2,
"sat_caption_l3": sat_l3,
"pair_id": entry["drone_name"],
}
@@ -248,5 +266,8 @@ def collate_gtauav_batch(
"caption_l1": [b["caption_l1"] for b in batch],
"caption_l2": [b["caption_l2"] for b in batch],
"caption_l3": [b["caption_l3"] for b in batch],
"sat_caption_l1": [b["sat_caption_l1"] for b in batch],
"sat_caption_l2": [b["sat_caption_l2"] for b in batch],
"sat_caption_l3": [b["sat_caption_l3"] for b in batch],
"pair_ids": [b["pair_id"] for b in batch],
}

6
src/losses/multi_infonce.py
View File

@@ -145,7 +145,8 @@ class InfoNCELoss(nn.Module):
weight_b2a=self.weight_g2q,
)
gate = embeddings.get("gate", 1.0)
gate_q = embeddings.get("gate_q", embeddings.get("gate", 1.0))
gate_g = embeddings.get("gate_g", 1.0)
if isinstance(tau, float):
tau_out = torch.tensor(tau, device=loss.device)
@@ -155,5 +156,6 @@ class InfoNCELoss(nn.Module):
return {
"total": loss,
"temperature": tau_out,
"gate": torch.tensor(gate, device=loss.device),
"gate_q": torch.tensor(gate_q, device=loss.device),
"gate_g": torch.tensor(gate_g, device=loss.device),
}

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,
)

29
src/training/train_gtauav.py
View File

@@ -61,7 +61,6 @@ class TrainConfigGTAUAV:
dino_web_path: str = _DINO_WEB
dino_sat_path: str = _DINO_SAT
lrsclip_path: str = _LRSCLIP
proj_dim: int = 512
init_gate: float = 0.7
baseline_mode: bool = False
@@ -169,6 +168,9 @@ def _evaluate(
caption_l1=batch["caption_l1"],
caption_l2=batch["caption_l2"],
caption_l3=batch["caption_l3"],
sat_caption_l1=batch["sat_caption_l1"],
sat_caption_l2=batch["sat_caption_l2"],
sat_caption_l3=batch["sat_caption_l3"],
)
all_query.append(embeddings["query"].cpu())
all_gallery.append(embeddings["gallery"].cpu())
@@ -193,7 +195,8 @@ def _evaluate(
hit = (top_k == targets.unsqueeze(1)).any(dim=1).float()
metrics[f"r@{k}_g2q"] = float(hit.mean().item())
metrics["gate"] = model.fusion.gate_value
metrics["gate_q"] = model.fusion_query.gate_value
metrics["gate_g"] = model.fusion_gallery.gate_value
return metrics
@@ -233,7 +236,6 @@ def train(cfg: TrainConfigGTAUAV) -> None:
dino_web_path=cfg.dino_web_path,
dino_sat_path=cfg.dino_sat_path,
lrsclip_path=cfg.lrsclip_path,
proj_dim=cfg.proj_dim,
init_gate=cfg.init_gate,
baseline_mode=cfg.baseline_mode,
device=cfg.device,
@@ -372,6 +374,9 @@ def train(cfg: TrainConfigGTAUAV) -> None:
caption_l1=batch["caption_l1"],
caption_l2=batch["caption_l2"],
caption_l3=batch["caption_l3"],
sat_caption_l1=batch["sat_caption_l1"],
sat_caption_l2=batch["sat_caption_l2"],
sat_caption_l3=batch["sat_caption_l3"],
)
# Loss in fp32 (learnable temperature gradient overflows in fp16).
loss_dict = loss_fn(
@@ -402,19 +407,21 @@ def train(cfg: TrainConfigGTAUAV) -> None:
pbar.set_postfix(
loss=f"{total_loss.item():.3f}",
tau=f"{loss_dict['temperature'].item():.4f}",
gate=f"{loss_dict['gate'].item():.3f}",
gq=f"{loss_dict['gate_q'].item():.3f}",
gg=f"{loss_dict['gate_g'].item():.3f}",
)
elapsed = time.time() - epoch_start
means = {k: v / max(n_batches, 1) for k, v in agg.items()}
LOGGER.info(
"📈 epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate=%.4f",
"📈 epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate_q=%.4f gate_g=%.4f",
epoch, elapsed,
optimizer.param_groups[0]["lr"],
means.get("total", 0.0),
means.get("temperature", 0.0),
means.get("gate", 1.0),
means.get("gate_q", 1.0),
means.get("gate_g", 1.0),
)
epoch_record: dict = {
@@ -428,12 +435,13 @@ def train(cfg: TrainConfigGTAUAV) -> None:
val_metrics = _evaluate(model, test_loader, cfg.device)
epoch_record["val"] = val_metrics
LOGGER.info(
"🎯 val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f gate=%.4f",
"🎯 val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f",
epoch,
val_metrics.get("r@1_q2g", 0.0),
val_metrics.get("r@5_q2g", 0.0),
val_metrics.get("r@10_q2g", 0.0),
val_metrics.get("gate", 1.0),
val_metrics.get("gate_q", 1.0),
val_metrics.get("gate_g", 1.0),
)
history.append(epoch_record)
@@ -469,11 +477,12 @@ def train(cfg: TrainConfigGTAUAV) -> None:
LOGGER.info("✅ Training complete. Report: %s", report_path)
LOGGER.info(
"📊 Final — R@1=%.4f R@5=%.4f R@10=%.4f gate=%.4f",
"📊 Final — R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f",
final_metrics.get("r@1_q2g", 0.0),
final_metrics.get("r@5_q2g", 0.0),
final_metrics.get("r@10_q2g", 0.0),
final_metrics.get("gate", 1.0),
final_metrics.get("gate_q", 1.0),
final_metrics.get("gate_g", 1.0),
)