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

64
README.md
View File

@@ -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

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