536 lines
20 KiB
Python
536 lines
20 KiB
Python
"""Цикл обучения dual-encoder модели на GTA-UAV (symmetric fusion-архитектура).
|
||
|
||
И дрон, и спутник имеют собственное текстовое описание; слияние
|
||
(картинка+текст) происходит СИММЕТРИЧНО на обеих сторонах, каждая —
|
||
со своим экземпляром GatedFusion (веса не общие, т.к. визуальные
|
||
домены различаются).
|
||
|
||
Использование:
|
||
python train.py \\
|
||
--data_root /path/to/GTA-UAV \\
|
||
--descriptions_path /path/to/descriptions.json \\
|
||
--text_levels level1 \\
|
||
--dgtrs_checkpoint /path/to/DGTRS-CLIP-ViT-B-16 \\
|
||
--stripnet_checkpoint /path/to/stripnet_small.pth \\
|
||
--epochs 50 --batch_size 64 --micro_batch_size 64 --bf16
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import os
|
||
import random
|
||
import sys
|
||
import time
|
||
from contextlib import nullcontext
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
from torch.optim import AdamW
|
||
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||
|
||
from src.models.dual_encoder import build_dual_encoder, get_trainable_params
|
||
from src.losses import InfoNCELoss
|
||
from src.metrics import compute_bidirectional_metrics
|
||
from src.data.gta_uav import build_dataloaders
|
||
from src.data.gta_uav_eval import build_multipos_eval
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s %(name)s %(levelname)s %(message)s",
|
||
)
|
||
LOGGER = logging.getLogger("cvgl.train")
|
||
|
||
|
||
def set_seed(seed: int) -> None:
|
||
"""Зафиксировать все источники случайности (протокол §0.4: seed=42 dev).
|
||
|
||
Покрывает python random, numpy, torch (CPU/CUDA) и порядок сэмплирования
|
||
в DataLoader/semi-positive выборе — для воспроизводимого equal-budget.
|
||
"""
|
||
random.seed(seed)
|
||
np.random.seed(seed)
|
||
torch.manual_seed(seed)
|
||
torch.cuda.manual_seed_all(seed)
|
||
os.environ["PYTHONHASHSEED"] = str(seed)
|
||
LOGGER.info("🎲 Seed fixed: %d", seed)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# GPU info
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def log_gpu_info(device: torch.device) -> None:
|
||
if device.type != "cuda":
|
||
return
|
||
name = torch.cuda.get_device_name(device)
|
||
total = torch.cuda.get_device_properties(device).total_memory / 1024**3
|
||
LOGGER.info("🖥️ GPU: %s (%.1f GB VRAM)", name, total)
|
||
cap = torch.cuda.get_device_capability(device)
|
||
bf16_ok = cap[0] >= 8
|
||
LOGGER.info(
|
||
" Compute capability: %d.%d | BF16: %s",
|
||
cap[0], cap[1], "✅ supported" if bf16_ok else "❌ use FP16 instead",
|
||
)
|
||
|
||
|
||
def log_vram_usage(prefix: str = "") -> None:
|
||
if not torch.cuda.is_available():
|
||
return
|
||
allocated = torch.cuda.memory_allocated() / 1024**3
|
||
reserved = torch.cuda.memory_reserved() / 1024**3
|
||
LOGGER.info(" %sVRAM: %.2f GB allocated / %.2f GB reserved", prefix, allocated, reserved)
|
||
|
||
|
||
def get_amp_context(use_bf16: bool, use_fp16: bool, device: torch.device):
|
||
if use_bf16:
|
||
return torch.autocast(device_type=device.type, dtype=torch.bfloat16)
|
||
elif use_fp16:
|
||
return torch.autocast(device_type=device.type, dtype=torch.float16)
|
||
else:
|
||
return nullcontext()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Evaluation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@torch.no_grad()
|
||
@torch.no_grad()
|
||
def _encode_loader(model, loader, encode_fn, device, amp_ctx) -> torch.Tensor:
|
||
"""Прогнать loader через encode_fn (encode_drone/encode_satellite)."""
|
||
embs = []
|
||
for batch in loader:
|
||
images = batch["image"].to(device)
|
||
tokens = batch["tokens"].to(device)
|
||
with amp_ctx:
|
||
emb = encode_fn(images, tokens)
|
||
embs.append(emb.float().cpu())
|
||
return torch.cat(embs, dim=0)
|
||
|
||
|
||
def evaluate(
|
||
model: nn.Module,
|
||
eval_data: dict,
|
||
device: torch.device,
|
||
amp_ctx,
|
||
) -> dict[str, float]:
|
||
"""Multi-positive retrieval eval (протокол §6.2).
|
||
|
||
Дрон-запросы и УНИКАЛЬНАЯ спутниковая галерея кодируются раздельно;
|
||
для каждого дрона учитываются ВСЕ его positive/semi-positive тайлы
|
||
(hit-if-any). Считаются оба направления: q2g (primary) и g2q.
|
||
|
||
Args:
|
||
eval_data: dict из build_multipos_eval (drone_loader, gallery_loader,
|
||
positives_q2g, positives_g2q).
|
||
"""
|
||
model.eval()
|
||
|
||
query_emb = _encode_loader(
|
||
model, eval_data["drone_loader"], model.encode_drone, device, amp_ctx,
|
||
)
|
||
gallery_emb = _encode_loader(
|
||
model, eval_data["gallery_loader"], model.encode_satellite, device, amp_ctx,
|
||
)
|
||
|
||
metrics = compute_bidirectional_metrics(
|
||
query_emb, gallery_emb, ks=(1, 5, 10),
|
||
positives_q2g=eval_data["positives_q2g"],
|
||
positives_g2q=eval_data["positives_g2q"],
|
||
return_ranks=True,
|
||
)
|
||
# Отделяем per-query массивы (0-indexed ранги первого попадания) от
|
||
# числовых метрик: они пойдут в .npy для bootstrap CI / paired-теста,
|
||
# а не в history.json (иначе JSON-дамп упадёт на np.ndarray).
|
||
per_query = {
|
||
"q2g_best_ranks": metrics.pop("q2g_best_ranks"),
|
||
"g2q_best_ranks": metrics.pop("g2q_best_ranks"),
|
||
}
|
||
# Primary-алиасы: recall@k / mAP = q2g (для best-model и collect_results).
|
||
for k in (1, 5, 10):
|
||
metrics[f"recall@{k}"] = metrics[f"q2g_recall@{k}"]
|
||
metrics["mAP"] = metrics["q2g_mAP"]
|
||
|
||
model.train()
|
||
return metrics, per_query
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Training loop
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def train_one_epoch(
|
||
model: nn.Module,
|
||
train_loader,
|
||
criterion: nn.Module,
|
||
optimizer: torch.optim.Optimizer,
|
||
device: torch.device,
|
||
epoch: int,
|
||
grad_accumulate_steps: int = 1,
|
||
max_grad_norm: float = 1.0,
|
||
amp_ctx=None,
|
||
scaler: torch.amp.GradScaler | None = None,
|
||
) -> dict[str, float]:
|
||
model.train()
|
||
if amp_ctx is None:
|
||
amp_ctx = nullcontext()
|
||
|
||
total_loss = 0.0
|
||
total_acc_d2s = 0.0
|
||
total_acc_s2d = 0.0
|
||
n_steps = 0
|
||
|
||
optimizer.zero_grad()
|
||
|
||
for batch_idx, batch in enumerate(train_loader):
|
||
drone_images = batch["drone_image"].to(device)
|
||
drone_tokens = batch["drone_tokens"].to(device)
|
||
satellite_images = batch["satellite_image"].to(device)
|
||
satellite_tokens = batch["satellite_tokens"].to(device)
|
||
|
||
with amp_ctx:
|
||
outputs = model(drone_images, drone_tokens, satellite_images, satellite_tokens)
|
||
logits = outputs["logits"]
|
||
loss_dict = criterion(logits)
|
||
loss = loss_dict["loss"]
|
||
scaled_loss = loss / grad_accumulate_steps
|
||
|
||
if scaler is not None:
|
||
scaler.scale(scaled_loss).backward()
|
||
else:
|
||
scaled_loss.backward()
|
||
|
||
total_loss += loss.item()
|
||
total_acc_d2s += loss_dict["acc_d2s"].item()
|
||
total_acc_s2d += loss_dict["acc_s2d"].item()
|
||
n_steps += 1
|
||
|
||
if (batch_idx + 1) % grad_accumulate_steps == 0:
|
||
if scaler is not None:
|
||
scaler.unscale_(optimizer)
|
||
torch.nn.utils.clip_grad_norm_(
|
||
[p for p in model.parameters() if p.requires_grad],
|
||
max_grad_norm,
|
||
)
|
||
scaler.step(optimizer)
|
||
scaler.update()
|
||
else:
|
||
torch.nn.utils.clip_grad_norm_(
|
||
[p for p in model.parameters() if p.requires_grad],
|
||
max_grad_norm,
|
||
)
|
||
optimizer.step()
|
||
optimizer.zero_grad()
|
||
|
||
if (batch_idx + 1) % 50 == 0:
|
||
LOGGER.info(
|
||
" [Epoch %d] Step %d/%d | loss=%.4f | acc_d2s=%.3f | "
|
||
"acc_s2d=%.3f | τ=%.4f",
|
||
epoch, batch_idx + 1, len(train_loader),
|
||
loss.item(),
|
||
loss_dict["acc_d2s"].item(),
|
||
loss_dict["acc_s2d"].item(),
|
||
outputs["temperature"].item(),
|
||
)
|
||
|
||
if len(train_loader) % grad_accumulate_steps != 0:
|
||
if scaler is not None:
|
||
scaler.unscale_(optimizer)
|
||
torch.nn.utils.clip_grad_norm_(
|
||
[p for p in model.parameters() if p.requires_grad],
|
||
max_grad_norm,
|
||
)
|
||
scaler.step(optimizer)
|
||
scaler.update()
|
||
else:
|
||
torch.nn.utils.clip_grad_norm_(
|
||
[p for p in model.parameters() if p.requires_grad],
|
||
max_grad_norm,
|
||
)
|
||
optimizer.step()
|
||
optimizer.zero_grad()
|
||
|
||
return {
|
||
"loss": total_loss / max(n_steps, 1),
|
||
"acc_d2s": total_acc_d2s / max(n_steps, 1),
|
||
"acc_s2d": total_acc_s2d / max(n_steps, 1),
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def main(args):
|
||
set_seed(args.seed)
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
LOGGER.info("🚀 Device: %s", device)
|
||
log_gpu_info(device)
|
||
|
||
exp_name = f"exp_{'-'.join(args.text_levels)}_ep{args.epochs}_bs{args.batch_size}"
|
||
output_dir = Path(args.output_dir) / exp_name
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
LOGGER.info("📁 Output: %s", output_dir)
|
||
|
||
config = vars(args)
|
||
with open(output_dir / "config.json", "w") as f:
|
||
json.dump(config, f, indent=2, default=str)
|
||
|
||
train_loader, _ = build_dataloaders(
|
||
data_root=args.data_root,
|
||
descriptions_path=args.descriptions_path,
|
||
text_levels=args.text_levels,
|
||
train_meta=args.train_meta,
|
||
test_meta=args.test_meta,
|
||
batch_size=args.micro_batch_size,
|
||
num_workers=args.num_workers,
|
||
image_size=args.image_size,
|
||
build_test=False, # оценка идёт через multi-positive eval ниже
|
||
mutually_exclusive=args.mutually_exclusive,
|
||
seed=args.seed,
|
||
)
|
||
|
||
# Multi-positive eval (протокол §6.2): уникальная галерея + positive-карты.
|
||
eval_data = build_multipos_eval(
|
||
data_root=args.data_root,
|
||
test_meta=args.test_meta,
|
||
descriptions_path=args.descriptions_path,
|
||
text_levels=args.text_levels,
|
||
image_size=args.image_size,
|
||
batch_size=args.micro_batch_size,
|
||
num_workers=args.num_workers,
|
||
)
|
||
|
||
model = build_dual_encoder(
|
||
dgtrs_checkpoint=args.dgtrs_checkpoint,
|
||
stripnet_checkpoint=args.stripnet_checkpoint,
|
||
fused_dim=args.fused_dim,
|
||
shared_dim=args.shared_dim,
|
||
freeze_text=True,
|
||
freeze_image_backbone=True,
|
||
inject_mona=args.inject_mona,
|
||
mona_bottleneck=args.mona_bottleneck,
|
||
device=str(device),
|
||
)
|
||
log_vram_usage("After model load: ")
|
||
|
||
if args.compile and hasattr(torch, "compile"):
|
||
LOGGER.info("⚡ Compiling model with torch.compile (mode=%s)", args.compile_mode)
|
||
model = torch.compile(model, mode=args.compile_mode)
|
||
|
||
use_bf16 = args.bf16 and device.type == "cuda"
|
||
use_fp16 = args.fp16 and device.type == "cuda" and not use_bf16
|
||
amp_ctx = get_amp_context(use_bf16, use_fp16, device)
|
||
scaler = torch.amp.GradScaler("cuda") if use_fp16 else None
|
||
|
||
precision_str = "BF16" if use_bf16 else ("FP16" if use_fp16 else "FP32")
|
||
LOGGER.info("🔢 Precision: %s", precision_str)
|
||
|
||
trainable_params = [p for p in model.parameters() if p.requires_grad]
|
||
optimizer = AdamW(
|
||
trainable_params,
|
||
lr=args.lr,
|
||
weight_decay=args.weight_decay,
|
||
betas=(0.9, 0.98),
|
||
)
|
||
|
||
# Linear-warmup + cosine (протокол §8.1 / §0.1). Warmup зажимается до
|
||
# epochs-1, чтобы на коротких прогонах (смок epochs=1) деградировать
|
||
# к чистому косинусу без ошибок.
|
||
warmup_epochs = max(0, min(args.warmup_epochs, args.epochs - 1))
|
||
cosine = CosineAnnealingLR(
|
||
optimizer,
|
||
T_max=args.epochs - warmup_epochs,
|
||
eta_min=args.lr * 0.01,
|
||
)
|
||
if warmup_epochs > 0:
|
||
warmup = LinearLR(
|
||
optimizer, start_factor=0.01, end_factor=1.0, total_iters=warmup_epochs,
|
||
)
|
||
scheduler = SequentialLR(
|
||
optimizer, schedulers=[warmup, cosine], milestones=[warmup_epochs],
|
||
)
|
||
LOGGER.info("📉 LR schedule: linear-warmup(%d) + cosine(%d)", warmup_epochs, args.epochs - warmup_epochs)
|
||
else:
|
||
scheduler = cosine
|
||
LOGGER.info("📉 LR schedule: cosine (%d epochs, no warmup)", args.epochs)
|
||
|
||
criterion = InfoNCELoss(label_smoothing=args.label_smoothing)
|
||
|
||
grad_accumulate_steps = max(1, args.batch_size // args.micro_batch_size)
|
||
LOGGER.info(
|
||
"⚙️ Effective batch=%d (micro=%d × accumulate=%d)",
|
||
args.batch_size, args.micro_batch_size, grad_accumulate_steps,
|
||
)
|
||
|
||
start_epoch = 1
|
||
if args.resume and (output_dir / "latest_model.pth").exists():
|
||
ckpt = torch.load(output_dir / "latest_model.pth", map_location=device)
|
||
model.load_state_dict(ckpt["model_state_dict"])
|
||
optimizer.load_state_dict(ckpt["optimizer_state_dict"])
|
||
if "scheduler_state_dict" in ckpt:
|
||
scheduler.load_state_dict(ckpt["scheduler_state_dict"])
|
||
start_epoch = ckpt["epoch"] + 1
|
||
LOGGER.info("🔄 Resumed from epoch %d", start_epoch - 1)
|
||
|
||
best_recall1 = 0.0
|
||
history = []
|
||
|
||
history_path = output_dir / "history.json"
|
||
if args.resume and history_path.exists():
|
||
with open(history_path) as f:
|
||
history = json.load(f)
|
||
|
||
log_vram_usage("Before training: ")
|
||
|
||
for epoch in range(start_epoch, args.epochs + 1):
|
||
t0 = time.time()
|
||
|
||
train_metrics = train_one_epoch(
|
||
model, train_loader, criterion, optimizer,
|
||
device, epoch,
|
||
grad_accumulate_steps=grad_accumulate_steps,
|
||
max_grad_norm=args.max_grad_norm,
|
||
amp_ctx=amp_ctx,
|
||
scaler=scaler,
|
||
)
|
||
|
||
scheduler.step()
|
||
|
||
if epoch % args.eval_every == 0 or epoch == args.epochs:
|
||
eval_metrics, per_query = evaluate(model, eval_data, device, amp_ctx)
|
||
# Per-query ранги последнего eval (перезапись) — источник для
|
||
# bootstrap CI / paired-теста между вариантами (протокол §8.3).
|
||
np.save(output_dir / "last_ranks_q2g.npy", per_query["q2g_best_ranks"])
|
||
np.save(output_dir / "last_ranks_g2q.npy", per_query["g2q_best_ranks"])
|
||
else:
|
||
eval_metrics, per_query = {}, None
|
||
|
||
elapsed = time.time() - t0
|
||
|
||
LOGGER.info(
|
||
"📈 Epoch %d/%d (%.0fs) | loss=%.4f | "
|
||
"R@1=%.3f R@5=%.3f R@10=%.3f | mAP=%.3f",
|
||
epoch, args.epochs, elapsed,
|
||
train_metrics["loss"],
|
||
eval_metrics.get("recall@1", 0),
|
||
eval_metrics.get("recall@5", 0),
|
||
eval_metrics.get("recall@10", 0),
|
||
eval_metrics.get("mAP", 0),
|
||
)
|
||
|
||
if epoch == 1:
|
||
log_vram_usage("After first epoch: ")
|
||
|
||
record = {
|
||
"epoch": epoch,
|
||
"lr": scheduler.get_last_lr()[0],
|
||
**{f"train_{k}": v for k, v in train_metrics.items()},
|
||
**{f"eval_{k}": v for k, v in eval_metrics.items()},
|
||
"elapsed_s": elapsed,
|
||
}
|
||
history.append(record)
|
||
|
||
recall1 = eval_metrics.get("recall@1", 0)
|
||
if recall1 > best_recall1:
|
||
best_recall1 = recall1
|
||
torch.save(
|
||
{
|
||
"epoch": epoch,
|
||
"model_state_dict": model.state_dict(),
|
||
"eval_metrics": eval_metrics,
|
||
"config": config,
|
||
},
|
||
output_dir / "best_model.pth",
|
||
)
|
||
# Ранги best-эпохи — для честного сравнения вариантов по лучшей
|
||
# точке (bootstrap CI / McNemar в compare-скрипте после обучения).
|
||
if per_query is not None:
|
||
np.save(output_dir / "best_ranks_q2g.npy", per_query["q2g_best_ranks"])
|
||
np.save(output_dir / "best_ranks_g2q.npy", per_query["g2q_best_ranks"])
|
||
LOGGER.info("💾 New best model (R@1=%.4f)", recall1)
|
||
|
||
torch.save(
|
||
{
|
||
"epoch": epoch,
|
||
"model_state_dict": model.state_dict(),
|
||
"optimizer_state_dict": optimizer.state_dict(),
|
||
"scheduler_state_dict": scheduler.state_dict(),
|
||
},
|
||
output_dir / "latest_model.pth",
|
||
)
|
||
|
||
with open(history_path, "w") as f:
|
||
json.dump(history, f, indent=2)
|
||
|
||
LOGGER.info("=" * 60)
|
||
LOGGER.info("🏁 Training complete. Best R@1: %.4f", best_recall1)
|
||
LOGGER.info("📁 Results: %s", output_dir)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def parse_args():
|
||
p = argparse.ArgumentParser(description="Train CVGL fusion dual-encoder on GTA-UAV")
|
||
|
||
# Data
|
||
p.add_argument("--data_root", type=str, required=True)
|
||
p.add_argument("--descriptions_path", type=str, required=True)
|
||
p.add_argument("--text_levels", nargs="+", default=["level1"])
|
||
p.add_argument("--train_meta", default="cross-area-drone2sate-train.json")
|
||
p.add_argument("--test_meta", default="cross-area-drone2sate-test.json")
|
||
p.add_argument("--image_size", type=int, default=256)
|
||
p.add_argument("--num_workers", type=int, default=8)
|
||
|
||
# Model
|
||
p.add_argument("--dgtrs_checkpoint", type=str, required=True)
|
||
p.add_argument("--stripnet_checkpoint", type=str, required=True)
|
||
p.add_argument("--fused_dim", type=int, default=512,
|
||
help="Размерность вектора после слияния картинки и текста")
|
||
p.add_argument("--shared_dim", type=int, default=512)
|
||
p.add_argument("--inject_mona", action="store_true", default=True)
|
||
p.add_argument("--mona_bottleneck", type=int, default=64)
|
||
p.add_argument(
|
||
"--mutually_exclusive", action=argparse.BooleanOptionalAction, default=True,
|
||
help="Батчи без общих позитивных тайлов (§6.4); --no-mutually_exclusive "
|
||
"→ обычный random shuffle.",
|
||
)
|
||
|
||
# Training
|
||
p.add_argument("--epochs", type=int, default=50)
|
||
p.add_argument("--batch_size", type=int, default=64)
|
||
p.add_argument("--micro_batch_size", type=int, default=64)
|
||
p.add_argument("--lr", type=float, default=1e-4)
|
||
p.add_argument("--weight_decay", type=float, default=0.01)
|
||
p.add_argument("--max_grad_norm", type=float, default=1.0)
|
||
p.add_argument("--label_smoothing", type=float, default=0.1)
|
||
p.add_argument("--eval_every", type=int, default=1)
|
||
|
||
# Performance
|
||
p.add_argument("--bf16", action="store_true", default=True)
|
||
p.add_argument("--fp16", action="store_true", default=False)
|
||
p.add_argument("--compile", action="store_true", default=False)
|
||
p.add_argument("--compile_mode", default="reduce-overhead",
|
||
choices=["default", "reduce-overhead", "max-autotune"])
|
||
|
||
# Resume / output
|
||
p.add_argument("--resume", action="store_true", default=False)
|
||
p.add_argument("--output_dir", type=str, default="outputs")
|
||
p.add_argument("--seed", type=int, default=42)
|
||
p.add_argument("--warmup_epochs", type=int, default=3,
|
||
help="Эпох линейного warmup перед косинусом (§8.1).")
|
||
|
||
return p.parse_args()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main(parse_args())
|