568 lines
18 KiB
Python
568 lines
18 KiB
Python
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 math
|
|
import time
|
|
import warnings
|
|
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 LambdaLR
|
|
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,
|
|
get_drone_train_transform,
|
|
get_satellite_train_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 = "meta/train_80.json"
|
|
_TEST_JSON = "meta/test_20.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
|
|
init_gate: float = 0.7
|
|
baseline_mode: bool = False
|
|
|
|
# Training.
|
|
resume_from: str | None = None # path to checkpoint for resuming
|
|
output_dir: str = "out/gtauav/with_text"
|
|
epochs: int = 20
|
|
batch_size: int = 8
|
|
num_workers: int = 4
|
|
learning_rate: float = 1e-4
|
|
text_lr_factor: float = 0.1 # text encoder LR = learning_rate * factor
|
|
weight_decay: float = 1e-4
|
|
grad_clip: float = 1.0
|
|
use_amp: bool = True
|
|
eval_every: int = 2
|
|
warmup_epochs: int = 2
|
|
seed: int = 42
|
|
device: str = "cuda"
|
|
|
|
# Loss.
|
|
tau_init: float = 0.07
|
|
label_smoothing: float = 0.1
|
|
weight_q2g: float = 0.6
|
|
weight_g2q: float = 0.4
|
|
learnable_temperature: bool = True
|
|
|
|
|
|
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)
|
|
|
|
|
|
def _build_param_groups(
|
|
model: AsymmetricEncoder,
|
|
lr: float,
|
|
text_lr_factor: float,
|
|
) -> list[dict]:
|
|
"""Build optimizer param groups with separate LR for text encoder."""
|
|
text_params = []
|
|
other_params = []
|
|
|
|
for name, param in model.named_parameters():
|
|
if not param.requires_grad:
|
|
continue
|
|
if "text_encoder" in name:
|
|
text_params.append(param)
|
|
else:
|
|
other_params.append(param)
|
|
|
|
groups = [{"params": other_params, "lr": lr}]
|
|
if text_params:
|
|
groups.append({"params": text_params, "lr": lr * text_lr_factor})
|
|
|
|
return groups
|
|
|
|
|
|
def _cosine_warmup_schedule(
|
|
warmup_steps: int,
|
|
total_steps: int,
|
|
) -> callable:
|
|
"""Cosine annealing with linear warmup."""
|
|
|
|
def lr_lambda(step: int) -> float:
|
|
if step < warmup_steps:
|
|
return step / max(warmup_steps, 1)
|
|
progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
|
|
return 0.5 * (1.0 + math.cos(math.pi * progress))
|
|
|
|
return lr_lambda
|
|
|
|
|
|
@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"],
|
|
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())
|
|
|
|
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_q"] = model.fusion_query.gate_value
|
|
metrics["gate_g"] = model.fusion_gallery.gate_value
|
|
return metrics
|
|
|
|
|
|
def _clear_vram() -> None:
|
|
"""Free VRAM from previous runs before starting."""
|
|
import gc
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.reset_peak_memory_stats()
|
|
allocated = torch.cuda.memory_allocated() / 1e9
|
|
LOGGER.info("🧹 VRAM cleared. Current usage: %.2f GB", allocated)
|
|
|
|
|
|
def train(cfg: TrainConfigGTAUAV) -> None:
|
|
"""Run full training loop."""
|
|
coloredlogs.install(
|
|
level="INFO",
|
|
logger=LOGGER,
|
|
fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
|
|
)
|
|
_clear_vram()
|
|
_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.
|
|
start_epoch = 0
|
|
resume_ckpt = None
|
|
|
|
if cfg.resume_from is not None:
|
|
LOGGER.info("🔄 Resuming from %s", cfg.resume_from)
|
|
model, resume_ckpt = AsymmetricEncoder.load_checkpoint(
|
|
cfg.resume_from,
|
|
dino_web_path=cfg.dino_web_path,
|
|
dino_sat_path=cfg.dino_sat_path,
|
|
lrsclip_path=cfg.lrsclip_path,
|
|
device=cfg.device,
|
|
)
|
|
start_epoch = resume_ckpt.get("epoch", -1) + 1
|
|
else:
|
|
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,
|
|
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,
|
|
label_smoothing=cfg.label_smoothing,
|
|
weight_q2g=cfg.weight_q2g,
|
|
weight_g2q=cfg.weight_g2q,
|
|
learnable_temperature=cfg.learnable_temperature,
|
|
)
|
|
LOGGER.info(
|
|
"🌡️ Temperature: %s (init=%.3f)",
|
|
"learnable" if cfg.learnable_temperature else "cosine schedule",
|
|
cfg.tau_init,
|
|
)
|
|
|
|
# Data — separate transforms for train (augmented) and eval (clean).
|
|
drone_train_tf = get_drone_train_transform(image_size=256)
|
|
sat_train_tf = get_satellite_train_transform(image_size=256)
|
|
eval_tf = get_dino_transform(image_size=256)
|
|
|
|
train_ds = GTAUAVDataset(
|
|
pair_json=cfg.train_json,
|
|
rgb_root=cfg.rgb_root,
|
|
caption_root=cfg.caption_root,
|
|
drone_transform=drone_train_tf,
|
|
sat_transform=sat_train_tf,
|
|
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=eval_tf,
|
|
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 — per-group LR (text encoder gets lower LR).
|
|
param_groups = _build_param_groups(model, cfg.learning_rate, cfg.text_lr_factor)
|
|
# Include loss temperature if learnable.
|
|
if cfg.learnable_temperature and loss_fn.logit_scale is not None:
|
|
param_groups[0]["params"].append(loss_fn.logit_scale)
|
|
|
|
optimizer = AdamW(param_groups, weight_decay=cfg.weight_decay)
|
|
|
|
lr_info = f"proj={cfg.learning_rate:.0e}"
|
|
if not cfg.baseline_mode:
|
|
lr_info += f" text={cfg.learning_rate * cfg.text_lr_factor:.0e}"
|
|
LOGGER.info("⚙️ Optimizer: AdamW LR: %s warmup=%d epochs", lr_info, cfg.warmup_epochs)
|
|
|
|
# Scheduler — cosine with linear warmup.
|
|
steps_per_epoch = len(train_loader)
|
|
total_steps = cfg.epochs * steps_per_epoch
|
|
warmup_steps = cfg.warmup_epochs * steps_per_epoch
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
|
|
scheduler = LambdaLR(
|
|
optimizer,
|
|
lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps),
|
|
last_epoch=-1,
|
|
)
|
|
scaler = GradScaler(enabled=cfg.use_amp)
|
|
|
|
# Restore optimizer/scheduler/loss state on resume.
|
|
if resume_ckpt is not None:
|
|
if "optimizer_state" in resume_ckpt:
|
|
optimizer.load_state_dict(resume_ckpt["optimizer_state"])
|
|
LOGGER.info("🔄 Optimizer state restored")
|
|
if "loss_state" in resume_ckpt:
|
|
loss_fn.load_state_dict(resume_ckpt["loss_state"])
|
|
LOGGER.info("🔄 Loss state restored (tau=%.4f)", loss_fn.current_temperature)
|
|
# Set scheduler last_epoch so it resumes at the correct LR.
|
|
scheduler.last_epoch = start_epoch * steps_per_epoch
|
|
LOGGER.info("🔄 Resuming from epoch %d", start_epoch)
|
|
|
|
history: list[dict] = []
|
|
|
|
LOGGER.info("🚀 Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch)
|
|
|
|
for epoch in range(start_epoch, cfg.epochs):
|
|
model.train()
|
|
epoch_start = time.time()
|
|
agg: dict[str, float] = {}
|
|
n_batches = 0
|
|
|
|
pbar = tqdm(
|
|
train_loader,
|
|
desc=f" 🏋️ Epoch {epoch + 1}/{cfg.epochs}",
|
|
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)
|
|
|
|
# Model forward in AMP (fp16 for DINOv3/DGTRS encoders).
|
|
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"],
|
|
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(
|
|
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()
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
|
|
scheduler.step()
|
|
|
|
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}",
|
|
tau=f"{loss_dict['temperature'].item():.4f}",
|
|
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_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_q", 1.0),
|
|
means.get("gate_g", 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_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_q", 1.0),
|
|
val_metrics.get("gate_g", 1.0),
|
|
)
|
|
|
|
history.append(epoch_record)
|
|
|
|
# Save checkpoint.
|
|
_atomic_save(
|
|
obj={
|
|
"epoch": epoch,
|
|
"model_state": model.state_dict(),
|
|
"optimizer_state": optimizer.state_dict(),
|
|
"loss_state": loss_fn.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_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_q", 1.0),
|
|
final_metrics.get("gate_g", 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(
|
|
"--resume", type=str, default=None,
|
|
help="Path to checkpoint to resume training from.",
|
|
)
|
|
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=8,
|
|
help="Batch size.",
|
|
)
|
|
parser.add_argument(
|
|
"--epochs", type=int, default=20,
|
|
help="Number of epochs.",
|
|
)
|
|
parser.add_argument(
|
|
"--lr", type=float, default=1e-4,
|
|
help="Learning rate for projections.",
|
|
)
|
|
parser.add_argument(
|
|
"--text-lr-factor", type=float, default=0.1,
|
|
help="Text encoder LR = lr * factor (default 0.1 = 10x lower).",
|
|
)
|
|
parser.add_argument(
|
|
"--warmup-epochs", type=int, default=2,
|
|
help="Linear warmup epochs.",
|
|
)
|
|
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.resume_from = args.resume
|
|
cfg.batch_size = args.batch_size
|
|
cfg.epochs = args.epochs
|
|
cfg.learning_rate = args.lr
|
|
cfg.text_lr_factor = args.text_lr_factor
|
|
cfg.warmup_epochs = args.warmup_epochs
|
|
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()
|