From 0c41c1f01792fd9a5566724933d40e8fb5fc5236 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Tue, 21 Apr 2026 19:01:30 +0300 Subject: [PATCH] Remove projections (1024 native), add satellite text, dual GatedFusion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- CLAUDE.md | 57 +++++++++------- README.md | 64 ++++++++++-------- src/datasets/gtauav_dataset.py | 25 ++++++- src/losses/multi_infonce.py | 6 +- src/models/asymmetric_encoder.py | 111 ++++++++++++++++--------------- src/training/train_gtauav.py | 29 +++++--- 6 files changed, 172 insertions(+), 120 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 4e9dbd3..c4400e5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -4,38 +4,45 @@ ``` QUERY BRANCH (drone + L1/L2/L3 captions): - drone_img [B,3,256,256] --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> CLS [B,1024] - | - proj_drone: Linear(1024,512) - | - L1 (overview) --> DGTRS-CLIP (248 tok) --> z₁ [B,768] --\ - L2 (full desc) --> DGTRS-CLIP (248 tok) --> z₂ [B,768] ---+-- cat --> [B,2304] - L3 (fingerprint) --> DGTRS-CLIP (248 tok) --> z₃ [B,768] --/ | - MLP(2304→768→512) - | - q = σ(α)·d_img + (1−σ(α))·d_txt GatedFusion - | - q̂ = q/‖q‖₂ --> query [B,512] + drone_img [B,3,256,256] --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> d_img [B,1024] + | + L1 --> DGTRS-CLIP (248 tok) --> z₁ [768] --\ | + L2 --> DGTRS-CLIP (248 tok) --> z₂ [768] ---+-- cat --> MLP(2304→1024→1024) --> d_txt [B,1024] + L3 --> DGTRS-CLIP (248 tok) --> z₃ [768] --/ | + | + q = σ(α_q)·d_img + (1−σ(α_q))·d_txt GatedFusion_q + | + q̂ = q/‖q‖₂ --> query [B,1024] -GALLERY BRANCH (satellite only): - sat_img [B,3,256,256] --> DINOv3 ViT-L/16 SAT-493M (frozen) --> CLS [B,1024] - | - proj_sat: Linear(1024,512) - | - ĝ = g/‖g‖₂ --> gallery [B,512] +GALLERY BRANCH (satellite + satellite captions): + sat_img [B,3,256,256] --> DINOv3 ViT-L/16 SAT-493M (frozen) --> s_img [B,1024] + | + sat_L1 --> DGTRS-CLIP --> z₁ --\ | + sat_L2 --> DGTRS-CLIP --> z₂ ---+-- cat --> MLP (shared) --> s_txt [B,1024] + 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) τ = 1/exp(logit_scale), learnable, clamped [0.01, 0.5], init=0.07 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) - **L1 overview:** первое предложение P1 — краткое описание land-cover (15-30 tok) - **L2 full:** полные P1 + P2 — inventory + spatial layout (100-200 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) - Код: `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 - Tokenizer: BPE SimpleTokenizer (248 tokens, vocab 49408) -### Trainable parameters: 10.9M из 733M (1.49%) -- proj_drone: Linear(1024,512) = ~524K -- proj_sat: Linear(1024,512) = ~524K -- TextFusionMLP: Linear(2304,768)+GELU+Linear(768,512) = ~2.2M -- gate alpha: 1 scalar +### Trainable parameters: 11.1M из 734M (1.51%) +- TextFusionMLP (shared): Linear(2304,1024)+GELU+Linear(1024,1024) = ~3.5M +- gate α_q: 1 scalar (query branch) +- gate α_g: 1 scalar (gallery branch) - logit_scale: 1 scalar (learnable temperature) - DGTRS partial unfreeze (last resblock + ln_final + text_projection): ~7.6M - DINOv3 x2 (303M each): frozen +- **Без projection layers** — retrieval space = DINOv3 native 1024-dim ### Optimizer & Scheduler - **AdamW** с per-group LR: projections lr=1e-4, text encoder lr=1e-5 diff --git a/README.md b/README.md index ad9bc0e..054a7b8 100644 --- a/README.md +++ b/README.md @@ -14,35 +14,39 @@ text fusion for drone-to-satellite image retrieval. ┌──────────────────────────── QUERY BRANCH ────────────────────────────┐ │ │ │ drone_img ──► DINOv3 ViT-L/16 LVD ──► CLS token │ -│ [B,3,256,256] (frozen, 303M) [B,1024] │ -│ │ │ -│ proj_drone: Linear(1024,512) │ -│ │ │ -│ d_img [B,512] │ +│ [B,3,256,256] (frozen, 303M) d_img [B,1024] │ │ │ │ │ L1 (overview) ──► DGTRS-CLIP ──► z₁ [B,768] ─┐ │ │ L2 (full desc) ──► DGTRS-CLIP ──► z₂ [B,768] ─┼─ cat ──► [B,2304]│ │ L3 (fingerprint) ──► DGTRS-CLIP ──► z₃ [B,768] ─┘ │ │ -│ (248 tokens, KPS pos. emb.) MLP(2304→768→512) │ -│ │ │ -│ d_txt [B,512] │ -│ │ │ -│ q = σ(α)·d_img + (1−σ(α))·d_txt GatedFusion │ +│ (248 tokens, KPS pos. emb.) MLP(2304→1024→1024) │ +│ │ │ +│ d_txt [B,1024] │ +│ │ │ +│ q = σ(α_q)·d_img + (1−σ(α_q))·d_txt GatedFusion_q │ │ │ │ -│ q̂ = q / ‖q‖₂ ──► query [B,512] │ +│ q̂ = q / ‖q‖₂ ──► query [B,1024] │ └───────────────────────────────────────────────────────────────────────┘ ┌──────────────────────────── GALLERY BRANCH ──────────────────────────┐ │ │ │ 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) @@ -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 (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₃] ) where [z₁ ; z₂ ; z₃] ∈ ℝ^(B×2304) — concatenation of three 768-dim embeddings - MLP: Linear(2304, 768) → GELU → Linear(768, 512) - z_text ∈ ℝ^(B×512) + MLP: Linear(2304, 1024) → GELU → Linear(1024, 1024) + 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 - d_img — projected drone image embedding [B, 512] - d_txt — fused text embedding [B, 512] + d_img, s_img — DINOv3 image embeddings [B, 1024] + d_txt, s_txt — fused text embeddings [B, 1024] + +For satellite images without captions: s_txt = None → g = s_img ``` ### Loss function @@ -113,7 +120,7 @@ Reported: R@1, R@5, R@10 for both q→g and g→q directions. ``` 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 - DGTRS text encoder (last resblock + ln_final + text_projection): 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 SAT (satellite) | 303M | 0 | frozen | | DGTRS-CLIP ViT-L-14 (text) | 124M | ~7.6M | last block + ln_final + text_projection | -| proj_drone | 524K | 524K | Linear(1024, 512) | -| proj_sat | 524K | 524K | Linear(1024, 512) | -| TextFusionMLP | 2.2M | 2.2M | Linear(2304,768) + GELU + Linear(768,512) | -| GatedFusion α | 1 | 1 | scalar | +| TextFusionMLP (shared) | 3.5M | 3.5M | Linear(2304,1024) + GELU + Linear(1024,1024) | +| GatedFusion α_q | 1 | 1 | query gate scalar | +| GatedFusion α_g | 1 | 1 | gallery gate scalar | | logit_scale | 1 | 1 | learnable temperature | -| **Total** | **733M** | **10.9M (1.49%)** | | +| **Total** | **734M** | **11.1M (1.51%)** | retrieval dim = 1024 | ## Experiments diff --git a/src/datasets/gtauav_dataset.py b/src/datasets/gtauav_dataset.py index ee129a0..5d8c1dc 100644 --- a/src/datasets/gtauav_dataset.py +++ b/src/datasets/gtauav_dataset.py @@ -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], } diff --git a/src/losses/multi_infonce.py b/src/losses/multi_infonce.py index ba70981..7749350 100644 --- a/src/losses/multi_infonce.py +++ b/src/losses/multi_infonce.py @@ -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), } diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 7a2785e..9ca6dd1 100644 --- a/src/models/asymmetric_encoder.py +++ b/src/models/asymmetric_encoder.py @@ -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, ) diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index cb26238..c946650 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -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), )