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

57
CLAUDE.md
View File

@@ -4,38 +4,45 @@
``` ```
QUERY BRANCH (drone + L1/L2/L3 captions): QUERY BRANCH (drone + L1/L2/L3 captions):
drone_img [B,3,256,256] --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> CLS [B,1024] drone_img [B,3,256,256] --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> d_img [B,1024]
| |
proj_drone: Linear(1024,512) L1 --> DGTRS-CLIP (248 tok) --> z₁ [768] --\ |
| L2 --> DGTRS-CLIP (248 tok) --> z₂ [768] ---+-- cat --> MLP(2304→1024→1024) --> d_txt [B,1024]
L1 (overview) --> DGTRS-CLIP (248 tok) --> z₁ [B,768] --\ L3 --> DGTRS-CLIP (248 tok) --> z₃ [768] --/ |
L2 (full desc) --> DGTRS-CLIP (248 tok) --> z₂ [B,768] ---+-- cat --> [B,2304] |
L3 (fingerprint) --> DGTRS-CLIP (248 tok) --> z₃ [B,768] --/ | q = σ(α_q)·d_img + (1−σ(α_q))·d_txt GatedFusion_q
MLP(2304→768→512) |
| q̂ = q/‖q‖₂ --> query [B,1024]
q = σ(α)·d_img + (1−σ(α))·d_txt GatedFusion
|
q̂ = q/‖q‖₂ --> query [B,512]
GALLERY BRANCH (satellite only): GALLERY BRANCH (satellite + satellite captions):
sat_img [B,3,256,256] --> DINOv3 ViT-L/16 SAT-493M (frozen) --> CLS [B,1024] sat_img [B,3,256,256] --> DINOv3 ViT-L/16 SAT-493M (frozen) --> s_img [B,1024]
| |
proj_sat: Linear(1024,512) sat_L1 --> DGTRS-CLIP --> z₁ --\ |
| sat_L2 --> DGTRS-CLIP --> z₂ ---+-- cat --> MLP (shared) --> s_txt [B,1024]
ĝ = g/‖g‖₂ --> gallery [B,512] sat_L3 --> DGTRS-CLIP --> z₃ --/ |
|
g = σ(α_g)·s_img + (1−σ(α_g))·s_txt GatedFusion_g
|
ĝ = g/‖g‖₂ --> gallery [B,1024]
Retrieval space: 1024-dim (DINOv3 native, без projection layers)
TextFusionMLP shared между query и gallery (одинаковый формат captions)
Для sat images без captions: s_txt=None → g = s_img (gate passthrough)
LOSS: L = 0.6·CE(q̂·ĝᵀ/τ, targets) + 0.4·CE(ĝ·q̂ᵀ/τ, targets) LOSS: L = 0.6·CE(q̂·ĝᵀ/τ, targets) + 0.4·CE(ĝ·q̂ᵀ/τ, targets)
τ = 1/exp(logit_scale), learnable, clamped [0.01, 0.5], init=0.07 τ = 1/exp(logit_scale), learnable, clamped [0.01, 0.5], init=0.07
label_smoothing=0.1 label_smoothing=0.1
BASELINE: σ(α) = 1.0, text branch disabled, DGTRS not loaded BASELINE: σ(α_q)=σ(α_g)=1.0, text disabled, DGTRS not loaded
``` ```
### Text hierarchy (L1/L2/L3) ### Text hierarchy (L1/L2/L3)
- **L1 overview:** первое предложение P1 — краткое описание land-cover (15-30 tok) - **L1 overview:** первое предложение P1 — краткое описание land-cover (15-30 tok)
- **L2 full:** полные P1 + P2 — inventory + spatial layout (100-200 tok) - **L2 full:** полные P1 + P2 — inventory + spatial layout (100-200 tok)
- **L3 fingerprint:** P3 — уникальные landmarks для matching (20-50 tok) - **L3 fingerprint:** P3 — уникальные landmarks для matching (20-50 tok)
- **Fusion:** z_text = MLP([z₁; z₂; z₃]) — concat 3×768 → Linear(2304,768) → GELU → Linear(768,512) - **Fusion:** z_text = MLP([z₁; z₂; z₃]) — concat 3×768 → Linear(2304,1024) → GELU → Linear(1024,1024)
- **Shared MLP** между query и gallery ветками (одинаковый формат captions)
- **Satellite captions:** 6,546 из 14,640 sat images имеют captions. Для остальных gate passthrough (g = s_img)
### Text encoder: DGTRS-CLIP (official architecture) ### Text encoder: DGTRS-CLIP (official architecture)
- Код: `src/models/dgtrs/` — из github.com/MitsuiChen14/DGTRS (Apache-2.0) - Код: `src/models/dgtrs/` — из github.com/MitsuiChen14/DGTRS (Apache-2.0)
@@ -43,14 +50,14 @@ BASELINE: σ(α) = 1.0, text branch disabled, DGTRS not loaded
- Transformer: sequence-first (LND), nn.MultiheadAttention, 12 layers - Transformer: sequence-first (LND), nn.MultiheadAttention, 12 layers
- Tokenizer: BPE SimpleTokenizer (248 tokens, vocab 49408) - Tokenizer: BPE SimpleTokenizer (248 tokens, vocab 49408)
### Trainable parameters: 10.9M из 733M (1.49%) ### Trainable parameters: 11.1M из 734M (1.51%)
- proj_drone: Linear(1024,512) = ~524K - TextFusionMLP (shared): Linear(2304,1024)+GELU+Linear(1024,1024) = ~3.5M
- proj_sat: Linear(1024,512) = ~524K - gate α_q: 1 scalar (query branch)
- TextFusionMLP: Linear(2304,768)+GELU+Linear(768,512) = ~2.2M - gate α_g: 1 scalar (gallery branch)
- gate alpha: 1 scalar
- logit_scale: 1 scalar (learnable temperature) - logit_scale: 1 scalar (learnable temperature)
- DGTRS partial unfreeze (last resblock + ln_final + text_projection): ~7.6M - DGTRS partial unfreeze (last resblock + ln_final + text_projection): ~7.6M
- DINOv3 x2 (303M each): frozen - DINOv3 x2 (303M each): frozen
- **Без projection layers** — retrieval space = DINOv3 native 1024-dim
### Optimizer & Scheduler ### Optimizer & Scheduler
- **AdamW** с per-group LR: projections lr=1e-4, text encoder lr=1e-5 - **AdamW** с per-group LR: projections lr=1e-4, text encoder lr=1e-5

64
README.md
View File

@@ -14,35 +14,39 @@ text fusion for drone-to-satellite image retrieval.
┌──────────────────────────── QUERY BRANCH ────────────────────────────┐ ┌──────────────────────────── QUERY BRANCH ────────────────────────────┐
│ │ │ │
│ drone_img ──► DINOv3 ViT-L/16 LVD ──► CLS token │ │ drone_img ──► DINOv3 ViT-L/16 LVD ──► CLS token │
│ [B,3,256,256] (frozen, 303M) [B,1024] │ │ [B,3,256,256] (frozen, 303M) d_img [B,1024] │
│ │ │
│ proj_drone: Linear(1024,512) │
│ │ │
│ d_img [B,512] │
│ │ │ │ │ │
│ L1 (overview) ──► DGTRS-CLIP ──► z₁ [B,768] ─┐ │ │ L1 (overview) ──► DGTRS-CLIP ──► z₁ [B,768] ─┐ │
│ L2 (full desc) ──► DGTRS-CLIP ──► z₂ [B,768] ─┼─ cat ──► [B,2304]│ │ L2 (full desc) ──► DGTRS-CLIP ──► z₂ [B,768] ─┼─ cat ──► [B,2304]│
│ L3 (fingerprint) ──► DGTRS-CLIP ──► z₃ [B,768] ─┘ │ │ │ L3 (fingerprint) ──► DGTRS-CLIP ──► z₃ [B,768] ─┘ │ │
│ (248 tokens, KPS pos. emb.) MLP(2304→768→512) │ │ (248 tokens, KPS pos. emb.) MLP(2304→1024→1024) │
│ │ │ │ │ │
│ d_txt [B,512] │ │ d_txt [B,1024] │
│ │ │ │ │ │
│ q = σ(α)·d_img + (1−σ(α))·d_txt GatedFusion │ │ q = σ(α_q)·d_img + (1−σ(α_q))·d_txt GatedFusion_q │
│ │ │ │ │ │
│ q̂ = q / ‖q‖₂ ──► query [B,512] │ │ q̂ = q / ‖q‖₂ ──► query [B,1024] │
└───────────────────────────────────────────────────────────────────────┘ └───────────────────────────────────────────────────────────────────────┘
┌──────────────────────────── GALLERY BRANCH ──────────────────────────┐ ┌──────────────────────────── GALLERY BRANCH ──────────────────────────┐
│ │ │ │
│ sat_img ──► DINOv3 ViT-L/16 SAT ──► CLS token │ │ sat_img ──► DINOv3 ViT-L/16 SAT ──► CLS token │
│ [B,3,256,256] (frozen, 303M) [B,1024] │ │ [B,3,256,256] (frozen, 303M) s_img [B,1024] │
│ │ │ │ │ │
│ proj_sat: Linear(1024,512) │ │ sat_L1 ──► DGTRS-CLIP ──► z₁ [768] ─┐ │
│ sat_L2 ──► DGTRS-CLIP ──► z₂ [768] ─┼─ cat ──► MLP ──► s_txt [1024]│
│ sat_L3 ──► DGTRS-CLIP ──► z₃ [768] ─┘ (shared MLP) │
│ │ │ │ │ │
│ ĝ = g / ‖g‖₂ ──► gallery [B,512] │ │ g = σ(α_g)·s_img + (1−σ(α_g))·s_txt GatedFusion_g │
│ │ │
│ ĝ = g / ‖g‖₂ ──► gallery [B,1024]│
└───────────────────────────────────────────────────────────────────────┘ └───────────────────────────────────────────────────────────────────────┘
BASELINE: σ(α) = 1.0 → q = d_img (text branch disabled, DGTRS not loaded) Retrieval space: 1024-dim (DINOv3 native, no projection layers)
TextFusionMLP shared between query and gallery branches
For sat images without captions: s_txt=None → g = s_img (gate passthrough)
BASELINE: σ(α) = 1.0 for both branches (text disabled, DGTRS not loaded)
``` ```
### Text hierarchy (L1 / L2 / L3) ### Text hierarchy (L1 / L2 / L3)
@@ -58,25 +62,28 @@ Each drone image has a VLM-generated caption (Qwen3-VL) split into 3 levels:
All three levels are encoded by a **single DGTRS-CLIP ViT-L-14** text encoder All three levels are encoded by a **single DGTRS-CLIP ViT-L-14** text encoder
(248-token context via KPS positional embedding, 768-dim output). (248-token context via KPS positional embedding, 768-dim output).
**Text fusion:** **Text fusion (shared MLP for both branches):**
``` ```
z_text = MLP( [z₁ ; z₂ ; z₃] ) z_text = MLP( [z₁ ; z₂ ; z₃] )
where [z₁ ; z₂ ; z₃] ∈ ℝ^(B×2304) — concatenation of three 768-dim embeddings where [z₁ ; z₂ ; z₃] ∈ ℝ^(B×2304) — concatenation of three 768-dim embeddings
MLP: Linear(2304, 768) → GELU → Linear(768, 512) MLP: Linear(2304, 1024) → GELU → Linear(1024, 1024)
z_text ∈ ℝ^(B×512) z_text ∈ ℝ^(B×1024)
``` ```
**Gated fusion:** **Gated fusion (separate gates for query and gallery):**
``` ```
q = σ(α) · d_img + (1 − σ(α)) · d_txt q = σ(α_q) · d_img + (1 − σ(α_q)) · d_txt (query branch)
g = σ(α_g) · s_img + (1 − σ(α_g)) · s_txt (gallery branch)
where α — learnable scalar in logit-space (init: σ(α) ≈ 0.7) where α_q, α_g — separate learnable scalars in logit-space (init: σ(α) ≈ 0.7)
σ — sigmoid function σ — sigmoid function
d_img — projected drone image embedding [B, 512] d_img, s_img — DINOv3 image embeddings [B, 1024]
d_txt — fused text embedding [B, 512] d_txt, s_txt — fused text embeddings [B, 1024]
For satellite images without captions: s_txt = None → g = s_img
``` ```
### Loss function ### Loss function
@@ -113,7 +120,7 @@ Reported: R@1, R@5, R@10 for both q→g and g→q directions.
``` ```
Optimizer: AdamW Optimizer: AdamW
- Projection heads (proj_drone, proj_sat, TextFusionMLP, gate α, logit_scale): - TextFusionMLP, gate α_q, gate α_g, logit_scale:
lr = 1e-4, weight_decay = 1e-4 lr = 1e-4, weight_decay = 1e-4
- DGTRS text encoder (last resblock + ln_final + text_projection): - DGTRS text encoder (last resblock + ln_final + text_projection):
lr = 1e-5 (10× lower, --text-lr-factor 0.1) lr = 1e-5 (10× lower, --text-lr-factor 0.1)
@@ -147,12 +154,11 @@ Mixed precision: AMP fp16 for model forward, fp32 for loss
| DINOv3 ViT-L/16 LVD (drone) | 303M | 0 | frozen | | DINOv3 ViT-L/16 LVD (drone) | 303M | 0 | frozen |
| DINOv3 ViT-L/16 SAT (satellite) | 303M | 0 | frozen | | DINOv3 ViT-L/16 SAT (satellite) | 303M | 0 | frozen |
| DGTRS-CLIP ViT-L-14 (text) | 124M | ~7.6M | last block + ln_final + text_projection | | DGTRS-CLIP ViT-L-14 (text) | 124M | ~7.6M | last block + ln_final + text_projection |
| proj_drone | 524K | 524K | Linear(1024, 512) | | TextFusionMLP (shared) | 3.5M | 3.5M | Linear(2304,1024) + GELU + Linear(1024,1024) |
| proj_sat | 524K | 524K | Linear(1024, 512) | | GatedFusion α_q | 1 | 1 | query gate scalar |
| TextFusionMLP | 2.2M | 2.2M | Linear(2304,768) + GELU + Linear(768,512) | | GatedFusion α_g | 1 | 1 | gallery gate scalar |
| GatedFusion α | 1 | 1 | scalar |
| logit_scale | 1 | 1 | learnable temperature | | logit_scale | 1 | 1 | learnable temperature |
| **Total** | **733M** | **10.9M (1.49%)** | | | **Total** | **734M** | **11.1M (1.51%)** | retrieval dim = 1024 |
## Experiments ## Experiments

25
src/datasets/gtauav_dataset.py
View File

@@ -176,13 +176,20 @@ class GTAUAVDataset(Dataset):
else: else:
continue # No match, skip. continue # No match, skip.
# Get captions. # Get drone captions.
cap_data = self.caption_index.get(drone_name) cap_data = self.caption_index.get(drone_name)
if cap_data is not None: if cap_data is not None:
l1, l2, l3 = _parse_caption_levels(cap_data["output"]) l1, l2, l3 = _parse_caption_levels(cap_data["output"])
else: else:
l1 = l2 = l3 = _EMPTY_CAPTION 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({ self.entries.append({
"drone_name": drone_name, "drone_name": drone_name,
"drone_dir": pair["drone_img_dir"], "drone_dir": pair["drone_img_dir"],
@@ -192,6 +199,7 @@ class GTAUAVDataset(Dataset):
"caption_l1": l1, "caption_l1": l1,
"caption_l2": l2, "caption_l2": l2,
"caption_l3": l3, "caption_l3": l3,
"sat_captions": sat_captions,
}) })
def _load_image(self, directory: str, filename: str, transform: Callable | None = None) -> torch.Tensor: 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) 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: if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
l1 = l2 = l3 = _EMPTY_CAPTION l1 = l2 = l3 = _EMPTY_CAPTION
else: else:
l1, l2, l3 = entry["caption_l1"], entry["caption_l2"], entry["caption_l3"] 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 { return {
"drone_img": drone_img, "drone_img": drone_img,
"sat_img": sat_img, "sat_img": sat_img,
"caption_l1": l1, "caption_l1": l1,
"caption_l2": l2, "caption_l2": l2,
"caption_l3": l3, "caption_l3": l3,
"sat_caption_l1": sat_l1,
"sat_caption_l2": sat_l2,
"sat_caption_l3": sat_l3,
"pair_id": entry["drone_name"], "pair_id": entry["drone_name"],
} }
@@ -248,5 +266,8 @@ def collate_gtauav_batch(
"caption_l1": [b["caption_l1"] for b in batch], "caption_l1": [b["caption_l1"] for b in batch],
"caption_l2": [b["caption_l2"] for b in batch], "caption_l2": [b["caption_l2"] for b in batch],
"caption_l3": [b["caption_l3"] 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], "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, 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): if isinstance(tau, float):
tau_out = torch.tensor(tau, device=loss.device) tau_out = torch.tensor(tau, device=loss.device)
@@ -155,5 +156,6 @@ class InfoNCELoss(nn.Module):
return { return {
"total": loss, "total": loss,
"temperature": tau_out, "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): class TextFusionMLP(nn.Module):
"""Fuse L1/L2/L3 text embeddings via concat + MLP. """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__( def __init__(
self, self,
text_dim: int = 768, text_dim: int = 768,
hidden_dim: int = 768, out_dim: int = 1024,
proj_dim: int = 512,
) -> None: ) -> None:
super().__init__() super().__init__()
self.mlp = nn.Sequential( self.mlp = nn.Sequential(
nn.Linear(3 * text_dim, hidden_dim), nn.Linear(3 * text_dim, out_dim),
nn.GELU(), nn.GELU(),
nn.Linear(hidden_dim, proj_dim), nn.Linear(out_dim, out_dim),
) )
def forward( def forward(
@@ -238,7 +237,7 @@ class TextFusionMLP(nn.Module):
z_l3: L3 fingerprint [B, text_dim]. z_l3: L3 fingerprint [B, text_dim].
Returns: Returns:
Fused text embedding [B, proj_dim]. Fused text embedding [B, out_dim].
""" """
cat = torch.cat([z_l1, z_l2, z_l3], dim=-1) cat = torch.cat([z_l1, z_l2, z_l3], dim=-1)
return self.mlp(cat) return self.mlp(cat)
@@ -249,18 +248,24 @@ class TextFusionMLP(nn.Module):
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
class AsymmetricEncoder(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 Query branch: DINOv3 LVD (drone) + text(L1/L2/L3) -> GatedFusion_q -> query [1024]
Gallery branch: DINOv3 SAT (satellite) -> gallery 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: Args:
dino_web_path: Path to DINOv3 LVD checkpoint (drone encoder). dino_web_path: Path to DINOv3 LVD checkpoint (drone encoder).
dino_sat_path: Path to DINOv3 SAT checkpoint (satellite encoder). dino_sat_path: Path to DINOv3 SAT checkpoint (satellite encoder).
lrsclip_path: Path to DGTRS-CLIP checkpoint (text encoder). lrsclip_path: Path to DGTRS-CLIP checkpoint (text encoder).
proj_dim: Shared projection dimension.
init_gate: Initial fusion gate (image weight). 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. 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_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors", dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt", lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
proj_dim: int = 512,
init_gate: float = 0.7, init_gate: float = 0.7,
baseline_mode: bool = False, baseline_mode: bool = False,
device: str = "cuda", device: str = "cuda",
) -> None: ) -> None:
super().__init__() super().__init__()
self.proj_dim = proj_dim self.embed_dim = self.DINO_DIM
self.baseline_mode = baseline_mode self.baseline_mode = baseline_mode
self.device = device self.device = device
@@ -296,24 +300,16 @@ class AsymmetricEncoder(nn.Module):
else: else:
self.text_encoder = None self.text_encoder = None
# Projection heads. # Shared text fusion MLP: 3×768 -> 1024 (same format for drone & sat captions).
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).
if not baseline_mode: if not baseline_mode:
self.text_fusion = TextFusionMLP( self.text_fusion = TextFusionMLP(
text_dim=self.TEXT_DIM, text_dim=self.TEXT_DIM,
hidden_dim=self.TEXT_DIM, out_dim=self.DINO_DIM,
proj_dim=proj_dim,
) )
# Gated fusion. # Separate gated fusion for query and gallery branches.
self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode) 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 @staticmethod
def _freeze(module: nn.Module) -> None: def _freeze(module: nn.Module) -> None:
@@ -354,8 +350,17 @@ class AsymmetricEncoder(nn.Module):
l1_texts: list[str], l1_texts: list[str],
l2_texts: list[str], l2_texts: list[str],
l3_texts: list[str], l3_texts: list[str],
) -> torch.Tensor: ) -> torch.Tensor | None:
"""Encode L1/L2/L3 captions and fuse. Returns [B, proj_dim].""" """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_l1 = self._encode_single_text(l1_texts)
z_l2 = self._encode_single_text(l2_texts) z_l2 = self._encode_single_text(l2_texts)
z_l3 = self._encode_single_text(l3_texts) z_l3 = self._encode_single_text(l3_texts)
@@ -374,45 +379,49 @@ class AsymmetricEncoder(nn.Module):
caption_l1: list[str] | None = None, caption_l1: list[str] | None = None,
caption_l2: list[str] | None = None, caption_l2: list[str] | None = None,
caption_l3: 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]: ) -> dict[str, torch.Tensor]:
"""Forward pass. """Forward pass.
Args: Args:
drone_img: Drone images [B, 3, 256, 256]. drone_img: Drone images [B, 3, 256, 256].
sat_img: Satellite images [B, 3, 256, 256]. sat_img: Satellite images [B, 3, 256, 256].
caption_l1: L1 overview captions. caption_l1/l2/l3: Drone L1/L2/L3 captions.
caption_l2: L2 full description captions. sat_caption_l1/l2/l3: Satellite L1/L2/L3 captions.
caption_l3: L3 fingerprint captions.
Returns: 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. # Image features (frozen DINOv3).
sat_feat = self.encode_satellite(sat_img)
gallery = self.proj_sat(sat_feat)
# Query: drone + optional text.
drone_feat = self.encode_drone(drone_img) drone_feat = self.encode_drone(drone_img)
drone_proj = self.proj_drone(drone_feat) sat_feat = self.encode_satellite(sat_img)
text_proj = None # Query branch: drone + drone text.
has_text = ( drone_text = None
caption_l1 is not None if (caption_l1 is not None and caption_l2 is not None
and caption_l2 is not None and caption_l3 is not None and not self.baseline_mode):
and caption_l3 is not None drone_text = self.encode_text_levels(caption_l1, caption_l2, caption_l3)
and not self.baseline_mode
)
if has_text:
text_proj = self.encode_text_levels(caption_l1, caption_l2, caption_l3)
query = self.fusion(drone_proj, text_proj) query = self.fusion_query(drone_feat, drone_text)
# Re-normalize after fusion.
query = F.normalize(query, dim=-1) 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 { return {
"query": query, "query": query,
"gallery": gallery, "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]: def trainable_parameters(self) -> list[nn.Parameter]:
@@ -425,7 +434,6 @@ class AsymmetricEncoder(nn.Module):
path.parent.mkdir(parents=True, exist_ok=True) path.parent.mkdir(parents=True, exist_ok=True)
ckpt = { ckpt = {
"model_state": self.state_dict(), "model_state": self.state_dict(),
"proj_dim": self.proj_dim,
"baseline_mode": self.baseline_mode, "baseline_mode": self.baseline_mode,
**extra, **extra,
} }
@@ -460,7 +468,6 @@ class AsymmetricEncoder(nn.Module):
dino_web_path=dino_web_path, dino_web_path=dino_web_path,
dino_sat_path=dino_sat_path, dino_sat_path=dino_sat_path,
lrsclip_path=lrsclip_path, lrsclip_path=lrsclip_path,
proj_dim=ckpt.get("proj_dim", 512),
baseline_mode=ckpt.get("baseline_mode", False), baseline_mode=ckpt.get("baseline_mode", False),
device=device, device=device,
) )

29
src/training/train_gtauav.py
View File

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