Add GTA-UAV experiment: asymmetric DINOv3 + LRSCLIP text encoder
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
425
src/training/train_gtauav.py
Normal file
425
src/training/train_gtauav.py
Normal 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()
|
||||
Reference in New Issue
Block a user