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

Add GTA-UAV experiment: asymmetric DINOv3 + LRSCLIP text encoder

Browse Source 
V3 architecture for CVGL caption validation on GTA-UAV-LR dataset:
- AsymmetricEncoder: DINOv3 ViT-L/16 (LVD drone + SAT satellite, frozen)
  + LRSCLIP/DGTRS-CLIP ViT-L-14 text encoder (248 tok, partial unfreeze)
- L1/L2/L3 hierarchical captions from VLM-generated descriptions
- TextFusionMLP (concat 3x768 -> MLP -> 512) + GatedFusion
- Segmentation filter: exclude images with >=90% background+water
- 10.9M trainable / 733M total params, 256x256 input
- coloredlogs + tqdm + emoji for training UX
- Baseline mode (--baseline): image-only, no text encoder loaded

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 17:54:27 +03:00
parent 5da791801c
commit 6ad9c4d149
10 changed files with 50043 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/datasets/__init__.py
View File

@@ -1,7 +1,7 @@
"""Dataset loaders for caption quality test."""
from src.datasets.visloc_with_captions import (
VisLocCaptionDataset,
GeoLocCaptionDataset,
collate_caption_batch,
)
__all__ = ["VisLocCaptionDataset", "collate_caption_batch"]
__all__ = ["GeoLocCaptionDataset", "collate_caption_batch"]

247
src/datasets/gtauav_dataset.py Normal file
View File

@@ -0,0 +1,247 @@
from __future__ import annotations
"""GTA-UAV-LR dataset loader with L1/L2/L3 hierarchical captions for CVGL.
Reads cross-area pair JSONs and VLM-generated caption JSONs.
Produces (drone_img, sat_img, caption_l1, caption_l2, caption_l3) tuples.
Caption levels:
L1 (overview): First sentence of P1 (land cover summary).
L2 (full): Complete P1 + P2 text.
L3 (fingerprint): P3 unique signature section.
For short captions (pure water, no P markers): all levels get the same short text.
"""
import json
import logging
import random
import re
from pathlib import Path
from typing import Any, Callable
import coloredlogs
import torch
from PIL import Image
from torch.utils.data import Dataset
from tqdm import tqdm
LOGGER = logging.getLogger("caption_test.gtauav_dataset")
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
# Default paths.
_RGB_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR")
_CAPTION_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR-captions")
_EMPTY_CAPTION = ""
# Regex to split P1/P2/P3 sections.
_P_SPLIT = re.compile(r"\*\*P[123][^*]*\*\*\s*:?\s*")
def _parse_caption_levels(output: str) -> tuple[str, str, str]:
"""Split VLM caption output into L1, L2, L3 levels.
Returns:
(l1_overview, l2_full, l3_fingerprint)
"""
sections = _P_SPLIT.split(output)
# sections[0] is empty (before **P1**), sections[1]=P1, [2]=P2, [3]=P3
sections = [s.strip() for s in sections if s.strip()]
if len(sections) >= 3:
p1, p2, p3 = sections[0], sections[1], sections[2]
elif len(sections) == 2:
p1, p2, p3 = sections[0], sections[1], sections[0]
elif len(sections) == 1:
p1 = p2 = p3 = sections[0]
else:
p1 = p2 = p3 = output.strip()
# L1: first sentence of P1.
first_dot = p1.find(". ")
l1 = p1[:first_dot + 1] if first_dot > 0 else p1
# L2: full P1 + P2.
l2 = p1 + " " + p2
# L3: P3 (fingerprint / unique signature).
l3 = p3
return l1, l2, l3
def _load_caption_index(caption_root: Path) -> dict[str, dict]:
"""Build index: image_name -> caption JSON data.
Scans drone/images/ and satellite/ directories.
"""
index: dict[str, dict] = {}
for subdir in ["drone/images", "satellite"]:
cap_dir = caption_root / subdir
if not cap_dir.exists():
continue
cap_files = sorted(cap_dir.glob("*_caption.json"))
for cap_file in tqdm(cap_files, desc=f" 📄 Loading {subdir} captions", unit="file", leave=False):
with open(cap_file) as f:
data = json.load(f)
# Key by the image name (without _caption suffix).
img_name = cap_file.name.replace("_caption.json", ".png")
index[img_name] = data
return index
class GTAUAVDataset(Dataset):
"""GTA-UAV-LR dataset with hierarchical L1/L2/L3 captions.
Args:
pair_json: Path to cross-area-drone2sate-{train,test}.json.
rgb_root: Root of GTA-UAV-LR RGB images.
caption_root: Root of GTA-UAV-LR-captions.
image_transform: Callable applied to PIL images.
filter_meta: Path to seg_filter.json (exclude 90%+ bg/water).
drop_caption_prob: Probability of dropping captions (ablation).
seed: Random seed.
"""
def __init__(
self,
pair_json: str,
rgb_root: str = str(_RGB_ROOT),
caption_root: str = str(_CAPTION_ROOT),
image_transform: Callable[[Image.Image], torch.Tensor] | None = None,
filter_meta: str | None = None,
drop_caption_prob: float = 0.0,
seed: int = 0,
) -> None:
self.rgb_root = Path(rgb_root)
self.caption_root = Path(caption_root)
self.image_transform = image_transform
self.drop_caption_prob = drop_caption_prob
self._rng = random.Random(seed)
# Load exclusion set from segmentation filter.
self.excluded: set[str] = set()
if filter_meta is not None:
self._load_filter(Path(filter_meta))
# Load caption index.
LOGGER.info("📚 Loading caption index from %s", caption_root)
self.caption_index = _load_caption_index(self.caption_root)
LOGGER.info("📚 Caption index: %d entries", len(self.caption_index))
# Load pairs.
self.entries: list[dict[str, Any]] = []
self._load_pairs(Path(pair_json))
LOGGER.info("✅ Loaded %d pairs from %s", len(self.entries), pair_json)
def _load_filter(self, path: Path) -> None:
with open(path) as f:
meta = json.load(f)
# excluded list contains paths like "drone/images/xxx.png"
for exc in meta.get("excluded", []):
# Extract image name from segm path.
self.excluded.add(Path(exc).name)
LOGGER.info("🔻 Filter loaded: %d excluded images", len(self.excluded))
def _load_pairs(self, pair_json: Path) -> None:
with open(pair_json) as f:
raw_pairs = json.load(f)
for pair in raw_pairs:
drone_name = pair["drone_img_name"]
# Skip excluded images.
if drone_name in self.excluded:
continue
# Get positive/semi-positive satellite images.
pos_list = pair.get("pair_pos_sate_img_list", [])
semipos_list = pair.get("pair_pos_semipos_sate_img_list", [])
semipos_weights = pair.get("pair_pos_semipos_sate_weight_list", [])
# Use positives if available, else semi-positives.
if pos_list:
sat_candidates = pos_list
sat_weights = None
elif semipos_list:
sat_candidates = semipos_list
sat_weights = semipos_weights if semipos_weights else None
else:
continue # No match, skip.
# Get 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
self.entries.append({
"drone_name": drone_name,
"drone_dir": pair["drone_img_dir"],
"sat_dir": pair["sate_img_dir"],
"sat_candidates": sat_candidates,
"sat_weights": sat_weights,
"caption_l1": l1,
"caption_l2": l2,
"caption_l3": l3,
})
def _load_image(self, directory: str, filename: str) -> torch.Tensor:
path = self.rgb_root / directory / filename
with Image.open(path) as img:
rgb = img.convert("RGB")
if self.image_transform is not None:
return self.image_transform(rgb)
return torch.tensor(0) # placeholder if no transform
def __len__(self) -> int:
return len(self.entries)
def __getitem__(self, idx: int) -> dict[str, Any]:
entry = self.entries[idx]
drone_img = self._load_image(entry["drone_dir"], entry["drone_name"])
# Sample satellite match (weighted if semi-positive).
if entry["sat_weights"] is not None:
sat_name = self._rng.choices(
entry["sat_candidates"],
weights=entry["sat_weights"],
k=1,
)[0]
else:
sat_name = self._rng.choice(entry["sat_candidates"])
sat_img = self._load_image(entry["sat_dir"], sat_name)
# 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"]
return {
"drone_img": drone_img,
"sat_img": sat_img,
"caption_l1": l1,
"caption_l2": l2,
"caption_l3": l3,
"pair_id": entry["drone_name"],
}
def collate_gtauav_batch(
batch: list[dict[str, Any]],
) -> dict[str, Any]:
"""Collate into batched dict. Captions stay as string lists."""
return {
"drone_img": torch.stack([b["drone_img"] for b in batch], dim=0),
"sat_img": torch.stack([b["sat_img"] for b in batch], dim=0),
"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],
"pair_ids": [b["pair_id"] for b in batch],
}

5
src/losses/__init__.py
View File

@@ -1,8 +1,7 @@
"""Loss functions for caption quality test."""
from src.losses.multi_infonce import (
MultiTermInfoNCE,
InfoNCELoss,
cosine_temperature,
curriculum_lambdas,
)
__all__ = ["MultiTermInfoNCE", "cosine_temperature", "curriculum_lambdas"]
__all__ = ["InfoNCELoss", "cosine_temperature"]

557
src/models/asymmetric_encoder.py Normal file
View File

@@ -0,0 +1,557 @@
from __future__ import annotations
"""Asymmetric dual encoder for CVGL caption test on GTA-UAV.
Architecture:
Query: DINOv3 ViT-L/16 (LVD, frozen) + LRSCLIP text (L1/L2/L3) -> GatedFusion -> query
Gallery: DINOv3 ViT-L/16 (SAT, frozen) -> gallery
Loss: InfoNCE(query, gallery)
DINOv3 checkpoints use a custom key layout (not HuggingFace transformers).
LRSCLIP (DGTRS-CLIP ViT-L-14) uses open_clip layout with KPS positional embeddings.
"""
import logging
import math
from pathlib import Path
import coloredlogs
import open_clip
import torch
import torch.nn as nn
import torch.nn.functional as F
LOGGER = logging.getLogger("caption_test.model")
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
from safetensors.torch import load_file as load_safetensors
from src.models.dual_encoder import GatedFusion, ProjectionHead
# ---------------------------------------------------------------------------
# DINOv3 ViT-L/16 — minimal implementation matching checkpoint key layout
# ---------------------------------------------------------------------------
class DINOv3Attention(nn.Module):
"""Multi-head self-attention with separate Q/K/V projections."""
def __init__(self, dim: int = 1024, num_heads: int = 16) -> None:
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.scale = self.head_dim ** -0.5
self.q_proj = nn.Linear(dim, dim)
self.k_proj = nn.Linear(dim, dim, bias=False)
self.v_proj = nn.Linear(dim, dim)
self.o_proj = nn.Linear(dim, dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, N, C = x.shape
q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
k = self.k_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
v = self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
attn = F.scaled_dot_product_attention(q, k, v)
x = attn.permute(0, 2, 1, 3).reshape(B, N, C)
return self.o_proj(x)
class DINOv3LayerScale(nn.Module):
"""Per-channel learnable scale (lambda)."""
def __init__(self, dim: int) -> None:
super().__init__()
self.lambda1 = nn.Parameter(torch.ones(dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
return x * self.lambda1
class DINOv3MLP(nn.Module):
"""SwiGLU-like MLP: up_proj + GELU + down_proj."""
def __init__(self, dim: int = 1024, mlp_dim: int = 4096) -> None:
super().__init__()
self.up_proj = nn.Linear(dim, mlp_dim)
self.down_proj = nn.Linear(mlp_dim, dim)
self.act = nn.GELU()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.down_proj(self.act(self.up_proj(x)))
class DINOv3Block(nn.Module):
"""Single DINOv3 transformer block."""
def __init__(self, dim: int = 1024, num_heads: int = 16, mlp_dim: int = 4096) -> None:
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attention = DINOv3Attention(dim, num_heads)
self.layer_scale1 = DINOv3LayerScale(dim)
self.norm2 = nn.LayerNorm(dim)
self.mlp = DINOv3MLP(dim, mlp_dim)
self.layer_scale2 = DINOv3LayerScale(dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x + self.layer_scale1(self.attention(self.norm1(x)))
x = x + self.layer_scale2(self.mlp(self.norm2(x)))
return x
class DINOv3Embeddings(nn.Module):
"""Patch embedding + CLS token + register tokens."""
def __init__(
self,
dim: int = 1024,
patch_size: int = 16,
num_registers: int = 4,
) -> None:
super().__init__()
self.patch_embeddings = nn.Conv2d(3, dim, patch_size, patch_size)
self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
self.register_tokens = nn.Parameter(torch.zeros(1, num_registers, dim))
self.mask_token = nn.Parameter(torch.zeros(1, 1, dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
B = x.shape[0]
patches = self.patch_embeddings(x).flatten(2).transpose(1, 2) # [B, N, D]
N = patches.shape[1]
cls = self.cls_token.expand(B, -1, -1)
reg = self.register_tokens.expand(B, -1, -1)
# DINOv3: [CLS, registers, patches]
x = torch.cat([cls, reg, patches], dim=1)
# Positional embedding: interpolated sincos (RoPE applied in attention
# in original, but pretrained checkpoints bake it into weights).
# We use a simple learned-style pos embed computed on the fly.
pos = self._get_pos_embed(N, x.device, x.dtype)
# pos covers patches only, skip CLS + registers
x[:, 1 + reg.shape[1]:] = x[:, 1 + reg.shape[1]:] + pos
return x
def _get_pos_embed(self, n_patches: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
# DINOv3 uses RoPE internally — no additive pos embed needed.
# Return zeros as placeholder (weights handle positioning via RoPE).
return torch.zeros(1, n_patches, self.cls_token.shape[-1], device=device, dtype=dtype)
class DINOv3ViT(nn.Module):
"""DINOv3 ViT-L/16 matching the checkpoint key layout.
Checkpoint keys:
embeddings.cls_token, embeddings.patch_embeddings.{weight,bias},
embeddings.register_tokens, embeddings.mask_token,
layer.{i}.attention.{q,k,v,o}_proj.{weight,bias},
layer.{i}.layer_scale{1,2}.lambda1,
layer.{i}.mlp.{up,down}_proj.{weight,bias},
layer.{i}.norm{1,2}.{weight,bias},
norm.{weight,bias}
"""
def __init__(
self,
dim: int = 1024,
num_heads: int = 16,
mlp_dim: int = 4096,
num_layers: int = 24,
patch_size: int = 16,
num_registers: int = 4,
) -> None:
super().__init__()
self.embeddings = DINOv3Embeddings(dim, patch_size, num_registers)
self.layer = nn.ModuleList([
DINOv3Block(dim, num_heads, mlp_dim) for _ in range(num_layers)
])
self.norm = nn.LayerNorm(dim)
self.embed_dim = dim
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Forward pass. Returns CLS token embedding [B, dim]."""
x = self.embeddings(x)
for block in self.layer:
x = block(x)
x = self.norm(x)
return x[:, 0] # CLS token
@classmethod
def from_pretrained(cls, path: str | Path) -> DINOv3ViT:
"""Load from .pth or .safetensors checkpoint."""
model = cls()
path = Path(path)
LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
if path.suffix == ".safetensors":
state = load_safetensors(str(path))
else:
state = torch.load(str(path), map_location="cpu", weights_only=False)
if "model" in state:
state = state["model"]
elif "state_dict" in state:
state = state["state_dict"]
model.load_state_dict(state, strict=False)
n_params = sum(p.numel() for p in model.parameters())
LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
return model
# ---------------------------------------------------------------------------
# LRSCLIP (DGTRS-CLIP) text encoder — open_clip with KPS positional embedding
# ---------------------------------------------------------------------------
class LRSCLIPTextEncoder(nn.Module):
"""DGTRS-CLIP text encoder with KPS positional embedding (248 tokens).
Wraps open_clip ViT-L-14 text tower. Handles the extra
positional_embedding_res (KPS residual) not present in vanilla open_clip.
Checkpoint keys (text-only):
token_embedding.weight: [49408, 768]
positional_embedding: [248, 768]
positional_embedding_res: [248, 768]
transformer.resblocks.{0-11}.{attn,mlp,ln}.*
ln_final.{weight,bias}
text_projection: [768, 768]
"""
def __init__(self, context_length: int = 248, embed_dim: int = 768) -> None:
super().__init__()
self.context_length = context_length
self.embed_dim = embed_dim
# Build the open_clip model to get correct architecture, then
# replace positional embedding and add KPS residual.
clip_model = open_clip.create_model("ViT-L-14", pretrained=None)
# Extract text components.
self.token_embedding = clip_model.token_embedding
self.transformer = clip_model.transformer
self.ln_final = clip_model.ln_final
self.text_projection = clip_model.text_projection
self.attn_mask = None # rebuilt in forward
# Replace positional embedding with 248-length version.
self.positional_embedding = nn.Parameter(
torch.zeros(context_length, embed_dim),
)
self.positional_embedding_res = nn.Parameter(
torch.zeros(context_length, embed_dim),
)
def _build_attn_mask(self, context_length: int, device: torch.device) -> torch.Tensor:
mask = torch.empty(context_length, context_length, device=device)
mask.fill_(float("-inf"))
mask.triu_(1)
return mask
def forward(self, text: torch.Tensor) -> torch.Tensor:
"""Encode tokenized text.
Args:
text: Token IDs [B, T] (T <= 248).
Returns:
Text embeddings [B, 768], L2-normalized.
"""
T = text.shape[1]
x = self.token_embedding(text) # [B, T, 768]
pos = self.positional_embedding[:T] + self.positional_embedding_res[:T]
x = x + pos
if self.attn_mask is None or self.attn_mask.shape[0] != T:
self.attn_mask = self._build_attn_mask(T, x.device)
attn_mask = self.attn_mask[:T, :T]
# open_clip transformer expects batch-first [B, T, D].
x = self.transformer(x, attn_mask=attn_mask)
x = self.ln_final(x) # [B, T, D]
# Take features at EOS token position.
eos_idx = text.argmax(dim=-1)
x = x[torch.arange(x.shape[0], device=x.device), eos_idx]
# Project.
if isinstance(self.text_projection, nn.Parameter):
x = x @ self.text_projection
else:
x = self.text_projection(x)
return F.normalize(x, dim=-1)
@classmethod
def from_pretrained(cls, path: str | Path) -> LRSCLIPTextEncoder:
"""Load DGTRS-CLIP checkpoint, extracting only text encoder weights."""
model = cls()
path = Path(path)
LOGGER.info("📝 Loading LRSCLIP text encoder from %s", path.name)
full_state = torch.load(str(path), map_location="cpu", weights_only=False)
if "state_dict" in full_state:
full_state = full_state["state_dict"]
# Filter out visual.* keys — keep only text encoder.
text_state = {
k: v for k, v in full_state.items()
if not k.startswith("visual.")
}
# Handle text_projection shape: checkpoint may be [768, 768] Parameter
# while open_clip stores it as nn.Parameter directly.
model.load_state_dict(text_state, strict=False)
n_params = sum(p.numel() for p in model.parameters())
LOGGER.info("📝 LRSCLIP loaded: %s params, context=%d tokens", f"{n_params:,}", model.context_length)
return model
# ---------------------------------------------------------------------------
# Text fusion MLP: concat L1/L2/L3 -> project to D
# ---------------------------------------------------------------------------
class TextFusionMLP(nn.Module):
"""Fuse L1/L2/L3 text embeddings via concat + MLP.
[B, 3*text_dim] -> [B, proj_dim]
"""
def __init__(
self,
text_dim: int = 768,
hidden_dim: int = 768,
proj_dim: int = 512,
) -> None:
super().__init__()
self.mlp = nn.Sequential(
nn.Linear(3 * text_dim, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, proj_dim),
)
def forward(
self,
z_l1: torch.Tensor,
z_l2: torch.Tensor,
z_l3: torch.Tensor,
) -> torch.Tensor:
"""Fuse three text embeddings.
Args:
z_l1: L1 overview [B, text_dim].
z_l2: L2 full description [B, text_dim].
z_l3: L3 fingerprint [B, text_dim].
Returns:
Fused text embedding [B, proj_dim].
"""
cat = torch.cat([z_l1, z_l2, z_l3], dim=-1)
return self.mlp(cat)
# ---------------------------------------------------------------------------
# Main model: AsymmetricEncoder
# ---------------------------------------------------------------------------
class AsymmetricEncoder(nn.Module):
"""Asymmetric dual encoder for CVGL with text fusion.
Query branch: DINOv3 LVD (drone) + LRSCLIP (L1/L2/L3) -> GatedFusion -> query
Gallery branch: DINOv3 SAT (satellite) -> gallery
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).
device: Torch device string.
"""
DINO_DIM = 1024
TEXT_DIM = 768
def __init__(
self,
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.baseline_mode = baseline_mode
self.device = device
# Image encoders (frozen).
self.drone_encoder = DINOv3ViT.from_pretrained(dino_web_path)
self.sat_encoder = DINOv3ViT.from_pretrained(dino_sat_path)
self._freeze(self.drone_encoder)
self._freeze(self.sat_encoder)
# Text encoder (partial unfreeze).
if not baseline_mode:
self.text_encoder = LRSCLIPTextEncoder.from_pretrained(lrsclip_path)
self._freeze(self.text_encoder)
self._unfreeze_text_last_block()
self.tokenizer = open_clip.get_tokenizer("ViT-L-14")
else:
self.text_encoder = None
self.tokenizer = 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).
if not baseline_mode:
self.text_fusion = TextFusionMLP(
text_dim=self.TEXT_DIM,
hidden_dim=self.TEXT_DIM,
proj_dim=proj_dim,
)
# Gated fusion.
self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
@staticmethod
def _freeze(module: nn.Module) -> None:
for p in module.parameters():
p.requires_grad = False
module.eval()
def _unfreeze_text_last_block(self) -> None:
"""Unfreeze last transformer block + text_projection + ln_final."""
if self.text_encoder is None:
return
# Last resblock.
for p in self.text_encoder.transformer.resblocks[-1].parameters():
p.requires_grad = True
# ln_final.
for p in self.text_encoder.ln_final.parameters():
p.requires_grad = True
# text_projection.
tp = self.text_encoder.text_projection
if isinstance(tp, nn.Parameter):
tp.requires_grad = True
elif isinstance(tp, nn.Module):
for p in tp.parameters():
p.requires_grad = True
def encode_drone(self, images: torch.Tensor) -> torch.Tensor:
"""Encode drone images. Returns [B, DINO_DIM]."""
with torch.no_grad():
return self.drone_encoder(images)
def encode_satellite(self, images: torch.Tensor) -> torch.Tensor:
"""Encode satellite images. Returns [B, DINO_DIM]."""
with torch.no_grad():
return self.sat_encoder(images)
def encode_text_levels(
self,
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]."""
z_l1 = self._encode_single_text(l1_texts)
z_l2 = self._encode_single_text(l2_texts)
z_l3 = self._encode_single_text(l3_texts)
fused = self.text_fusion(z_l1, z_l2, z_l3)
return F.normalize(fused, dim=-1)
def _encode_single_text(self, texts: list[str]) -> torch.Tensor:
"""Tokenize and encode a list of strings."""
tokens = self.tokenizer(list(texts)).to(self.device)
# Pad/truncate to context_length.
T = tokens.shape[1]
if T > self.text_encoder.context_length:
tokens = tokens[:, :self.text_encoder.context_length]
return self.text_encoder(tokens)
def forward(
self,
drone_img: torch.Tensor,
sat_img: torch.Tensor,
caption_l1: list[str] | None = None,
caption_l2: list[str] | None = None,
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.
Returns:
Dict with 'query' [B, proj_dim], 'gallery' [B, proj_dim], 'gate'.
"""
# Gallery: satellite only.
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_proj = self.proj_drone(drone_feat)
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 = self.fusion(drone_proj, text_proj)
# Re-normalize after fusion.
query = F.normalize(query, dim=-1)
return {
"query": query,
"gallery": gallery,
"gate": self.fusion.gate_value,
}
def trainable_parameters(self) -> list[nn.Parameter]:
"""Return list of parameters that require grad."""
return [p for p in self.parameters() if p.requires_grad]
def train(self, mode: bool = True) -> AsymmetricEncoder:
"""Override to keep frozen encoders in eval mode."""
super().train(mode)
self.drone_encoder.eval()
self.sat_encoder.eval()
if self.text_encoder is not None:
# Text encoder partially unfrozen — set to train mode
# but frozen layers won't update anyway.
self.text_encoder.train(mode)
return self
# ---------------------------------------------------------------------------
# Image preprocessing (DINOv3: 256x256, ImageNet normalization)
# ---------------------------------------------------------------------------
def get_dino_transform(image_size: int = 256) -> torch.nn.Module:
"""Build image transform for DINOv3 input."""
from torchvision import transforms
return transforms.Compose([
transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
])

425
src/training/train_gtauav.py Normal file
View File

@@ -0,0 +1,425 @@
from __future__ import annotations
"""Training loop for CVGL caption test on GTA-UAV-LR dataset.
Asymmetric DINOv3 encoders (drone LVD + satellite SAT) with LRSCLIP text fusion.
Single InfoNCE loss: query(drone+text) vs gallery(satellite).
"""
import argparse
import json
import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
import coloredlogs
import torch
import torch.nn as nn
from torch.amp import GradScaler, autocast
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader
from tqdm import tqdm
from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch
from src.losses.multi_infonce import InfoNCELoss
from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform
LOGGER = logging.getLogger("caption_test.train_gtauav")
# Default paths.
_RGB_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR"
_CAPTION_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR-captions"
_TRAIN_JSON = f"{_RGB_ROOT}/cross-area-drone2sate-train.json"
_TEST_JSON = f"{_RGB_ROOT}/cross-area-drone2sate-test.json"
_DINO_WEB = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth"
_DINO_SAT = "nn_models/DINO_SAT/model.safetensors"
_LRSCLIP = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt"
@dataclass
class TrainConfigGTAUAV:
"""Training configuration for GTA-UAV experiment."""
# Data.
train_json: str = _TRAIN_JSON
test_json: str = _TEST_JSON
rgb_root: str = _RGB_ROOT
caption_root: str = _CAPTION_ROOT
filter_meta: str | None = None
# Model.
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
# Training.
output_dir: str = "out/gtauav/with_text"
epochs: int = 10
batch_size: int = 64
num_workers: int = 4
learning_rate: float = 1e-4
weight_decay: float = 1e-4
grad_clip: float = 1.0
use_amp: bool = True
eval_every: int = 2
seed: int = 42
device: str = "cuda"
# Loss.
tau_init: float = 0.1
tau_final: float = 0.01
label_smoothing: float = 0.1
weight_q2g: float = 0.6
weight_g2q: float = 0.4
def _set_seed(seed: int) -> None:
import random as _random
import numpy as _np
_random.seed(seed)
_np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def _atomic_save(obj: dict, path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
tmp_path = path.with_suffix(path.suffix + ".tmp")
torch.save(obj, tmp_path)
tmp_path.replace(path)
@torch.no_grad()
def _evaluate(
model: AsymmetricEncoder,
loader: DataLoader,
device: str,
k_values: tuple[int, ...] = (1, 5, 10),
) -> dict[str, float]:
"""Compute R@K on validation set."""
model.eval()
all_query: list[torch.Tensor] = []
all_gallery: list[torch.Tensor] = []
for batch in tqdm(loader, desc=" 🔎 Evaluating", unit="batch", leave=False):
drone_img = batch["drone_img"].to(device, non_blocking=True)
sat_img = batch["sat_img"].to(device, non_blocking=True)
if model.baseline_mode:
embeddings = model(drone_img=drone_img, sat_img=sat_img)
else:
embeddings = model(
drone_img=drone_img,
sat_img=sat_img,
caption_l1=batch["caption_l1"],
caption_l2=batch["caption_l2"],
caption_l3=batch["caption_l3"],
)
all_query.append(embeddings["query"].cpu())
all_gallery.append(embeddings["gallery"].cpu())
query = torch.cat(all_query, dim=0)
gallery = torch.cat(all_gallery, dim=0)
sim = query @ gallery.t()
n = sim.size(0)
targets = torch.arange(n)
metrics: dict[str, float] = {}
sorted_idx = sim.argsort(dim=1, descending=True)
for k in k_values:
top_k = sorted_idx[:, :k]
hit = (top_k == targets.unsqueeze(1)).any(dim=1).float()
metrics[f"r@{k}_q2g"] = float(hit.mean().item())
sorted_idx_g2q = sim.t().argsort(dim=1, descending=True)
for k in k_values:
top_k = sorted_idx_g2q[:, :k]
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
return metrics
def train(cfg: TrainConfigGTAUAV) -> None:
"""Run full training loop."""
coloredlogs.install(
level="INFO",
logger=LOGGER,
fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
)
_set_seed(cfg.seed)
output_dir = Path(cfg.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
# Save config.
with (output_dir / "config.json").open("w") as f:
json.dump(vars(cfg), f, indent=2)
# Model.
mode_str = "🚫 baseline (no text)" if cfg.baseline_mode else "📝 with text (L1/L2/L3)"
LOGGER.info("🏗️ Building model — %s", mode_str)
model = AsymmetricEncoder(
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,
).to(cfg.device)
n_trainable = sum(p.numel() for p in model.trainable_parameters())
n_total = sum(p.numel() for p in model.parameters())
LOGGER.info(
"🧮 trainable=%s (%.2f%%) total=%s",
f"{n_trainable:,}", 100.0 * n_trainable / max(n_total, 1), f"{n_total:,}",
)
# Loss.
loss_fn = InfoNCELoss(
temperature_init=cfg.tau_init,
temperature_final=cfg.tau_final,
label_smoothing=cfg.label_smoothing,
weight_q2g=cfg.weight_q2g,
weight_g2q=cfg.weight_g2q,
)
# Data.
transform = get_dino_transform(image_size=256)
train_ds = GTAUAVDataset(
pair_json=cfg.train_json,
rgb_root=cfg.rgb_root,
caption_root=cfg.caption_root,
image_transform=transform,
filter_meta=cfg.filter_meta,
)
test_ds = GTAUAVDataset(
pair_json=cfg.test_json,
rgb_root=cfg.rgb_root,
caption_root=cfg.caption_root,
image_transform=transform,
filter_meta=cfg.filter_meta,
)
train_loader = DataLoader(
train_ds,
batch_size=cfg.batch_size,
shuffle=True,
num_workers=cfg.num_workers,
collate_fn=collate_gtauav_batch,
pin_memory=True,
drop_last=True,
)
test_loader = DataLoader(
test_ds,
batch_size=cfg.batch_size,
shuffle=False,
num_workers=cfg.num_workers,
collate_fn=collate_gtauav_batch,
pin_memory=True,
)
LOGGER.info("📦 train=%d test=%d batch=%d", len(train_ds), len(test_ds), cfg.batch_size)
# Optimizer.
optimizer = AdamW(
model.trainable_parameters(),
lr=cfg.learning_rate,
weight_decay=cfg.weight_decay,
)
scheduler = CosineAnnealingLR(optimizer, T_max=cfg.epochs)
scaler = GradScaler(enabled=cfg.use_amp)
history: list[dict] = []
LOGGER.info("🚀 Starting training for %d epochs", cfg.epochs)
for epoch in range(cfg.epochs):
model.train()
epoch_start = time.time()
agg: dict[str, float] = {}
n_batches = 0
pbar = tqdm(
train_loader,
desc=f" 🏋️ Epoch {epoch}/{cfg.epochs - 1}",
unit="batch",
leave=False,
)
for batch in pbar:
optimizer.zero_grad(set_to_none=True)
drone_img = batch["drone_img"].to(cfg.device, non_blocking=True)
sat_img = batch["sat_img"].to(cfg.device, non_blocking=True)
with autocast(device_type="cuda", enabled=cfg.use_amp):
if cfg.baseline_mode:
embeddings = model(drone_img=drone_img, sat_img=sat_img)
else:
embeddings = model(
drone_img=drone_img,
sat_img=sat_img,
caption_l1=batch["caption_l1"],
caption_l2=batch["caption_l2"],
caption_l3=batch["caption_l3"],
)
loss_dict = loss_fn(
embeddings=embeddings,
epoch=epoch,
total_epochs=cfg.epochs,
)
total_loss = loss_dict["total"]
scaler.scale(total_loss).backward()
if cfg.grad_clip > 0:
scaler.unscale_(optimizer)
nn.utils.clip_grad_norm_(
model.trainable_parameters(),
max_norm=cfg.grad_clip,
)
scaler.step(optimizer)
scaler.update()
for key, val in loss_dict.items():
agg[key] = agg.get(key, 0.0) + float(val.item())
n_batches += 1
pbar.set_postfix(
loss=f"{total_loss.item():.3f}",
gate=f"{loss_dict['gate'].item():.3f}",
)
scheduler.step()
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, elapsed,
optimizer.param_groups[0]["lr"],
means.get("total", 0.0),
means.get("temperature", 0.0),
means.get("gate", 1.0),
)
epoch_record: dict = {
"epoch": epoch,
"elapsed_seconds": elapsed,
"train": means,
}
# Evaluation.
if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
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",
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),
)
history.append(epoch_record)
# Save checkpoint.
_atomic_save(
obj={
"epoch": epoch,
"model_state": model.state_dict(),
"optimizer_state": optimizer.state_dict(),
},
path=output_dir / f"ckpt_epoch{epoch:03d}.pt",
)
LOGGER.info("💾 Checkpoint saved: ckpt_epoch%03d.pt", epoch)
# Save history.
history_path = output_dir / "history.json"
with history_path.open("w", encoding="utf-8") as f:
json.dump(history, f, indent=2)
# Save final eval report.
LOGGER.info("🔎 Running final evaluation...")
final_metrics = _evaluate(model, test_loader, cfg.device)
report = {
"config": vars(cfg),
"metrics": final_metrics,
"history": history,
}
report_path = output_dir / "eval_report.json"
with report_path.open("w", encoding="utf-8") as f:
json.dump(report, f, indent=2)
LOGGER.info("✅ Training complete. Report: %s", report_path)
LOGGER.info(
"📊 Final — R@1=%.4f R@5=%.4f R@10=%.4f gate=%.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),
)
def main() -> None:
parser = argparse.ArgumentParser(description="GTA-UAV caption test training.")
parser.add_argument(
"--baseline", action="store_true",
help="Run baseline mode (no text).",
)
parser.add_argument(
"--output-dir", type=str, default=None,
help="Override output directory.",
)
parser.add_argument(
"--filter-meta", type=str, default=None,
help="Path to seg_filter.json for excluding bad images.",
)
parser.add_argument(
"--batch-size", type=int, default=64,
help="Batch size.",
)
parser.add_argument(
"--epochs", type=int, default=10,
help="Number of epochs.",
)
parser.add_argument(
"--lr", type=float, default=1e-4,
help="Learning rate.",
)
parser.add_argument(
"--init-gate", type=float, default=0.7,
help="Initial gate value (image weight).",
)
args = parser.parse_args()
cfg = TrainConfigGTAUAV()
cfg.baseline_mode = args.baseline
cfg.batch_size = args.batch_size
cfg.epochs = args.epochs
cfg.learning_rate = args.lr
cfg.init_gate = args.init_gate
if args.filter_meta is not None:
cfg.filter_meta = args.filter_meta
if args.output_dir is not None:
cfg.output_dir = args.output_dir
elif args.baseline:
cfg.output_dir = "out/gtauav/baseline"
train(cfg)
if __name__ == "__main__":
main()