iait0/cvgl_experiments
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Full fix

Browse Source
This commit is contained in:
pikaliov
2026-07-07 18:47:38 +03:00
parent 3e95f4f618
commit 05355c6a7b
9 changed files with 599 additions and 337 deletions
Download Patch File Download Diff File Expand all files Collapse all files

179
train.py
View File

@@ -1,19 +1,10 @@
"""Цикл обучения dual-encoder модели на GTA-UAV (symmetric fusion-архитектура).
ВАЖНОЕ ИЗМЕНЕНИЕ: модель теперь принимает ЧЕТЫРЕ входа вместо трёх:
drone_images, drone_tokens → сливаются в drone_emb (TextFusionMLP)
satellite_images, satellite_tokens → сливаются в satellite_emb (TextFusionMLP)
drone_emb и satellite_emb сравниваются между собой → cosine similarity → InfoNCE
И дрон, и спутник имеют собственное текстовое описание; слияние
(картинка+текст) происходит СИММЕТРИЧНО на обеих сторонах, каждая —
со своим экземпляром TextFusionMLP (веса не общие, т.к. визуальные
со своим экземпляром GatedFusion (веса не общие, т.к. визуальные
домены различаются).
Оптимизировано под RTX 4090 (24 GB VRAM, Ada Lovelace): BF16 AMP,
micro_batch=64 по умолчанию (effective batch = micro_batch при отсутствии
gradient accumulation).
Использование:
python train.py \\
--data_root /path/to/GTA-UAV \\
@@ -28,22 +19,26 @@ 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
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_retrieval_metrics, format_metrics
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,
@@ -52,6 +47,20 @@ logging.basicConfig(
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
# ---------------------------------------------------------------------------
@@ -92,40 +101,64 @@ def get_amp_context(use_bf16: bool, use_fp16: bool, device: torch.device):
# ---------------------------------------------------------------------------
@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,
test_loader,
eval_data: dict,
device: torch.device,
amp_ctx,
) -> dict[str, float]:
"""Прогнать test set: drone (fused) vs satellite (fused) — симметрично."""
"""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()
all_drone_emb = []
all_satellite_emb = []
for batch in test_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:
drone_emb = model.encode_drone(drone_images, drone_tokens)
satellite_emb = model.encode_satellite(satellite_images, satellite_tokens)
all_drone_emb.append(drone_emb.float().cpu())
all_satellite_emb.append(satellite_emb.float().cpu())
all_drone_emb = torch.cat(all_drone_emb, dim=0)
all_satellite_emb = torch.cat(all_satellite_emb, dim=0)
metrics = compute_retrieval_metrics(
all_drone_emb, all_satellite_emb, ks=[1, 5, 10],
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
return metrics, per_query
# ---------------------------------------------------------------------------
@@ -149,8 +182,8 @@ def train_one_epoch(
amp_ctx = nullcontext()
total_loss = 0.0
total_acc_t2i = 0.0
total_acc_i2t = 0.0
total_acc_d2s = 0.0
total_acc_s2d = 0.0
n_steps = 0
optimizer.zero_grad()
@@ -174,8 +207,8 @@ def train_one_epoch(
scaled_loss.backward()
total_loss += loss.item()
total_acc_t2i += loss_dict["acc_t2i"].item()
total_acc_i2t += loss_dict["acc_i2t"].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:
@@ -201,8 +234,8 @@ def train_one_epoch(
"acc_s2d=%.3f | τ=%.4f",
epoch, batch_idx + 1, len(train_loader),
loss.item(),
loss_dict["acc_t2i"].item(),
loss_dict["acc_i2t"].item(),
loss_dict["acc_d2s"].item(),
loss_dict["acc_s2d"].item(),
outputs["temperature"].item(),
)
@@ -225,8 +258,8 @@ def train_one_epoch(
return {
"loss": total_loss / max(n_steps, 1),
"acc_t2i": total_acc_t2i / max(n_steps, 1),
"acc_i2t": total_acc_i2t / max(n_steps, 1),
"acc_d2s": total_acc_d2s / max(n_steps, 1),
"acc_s2d": total_acc_s2d / max(n_steps, 1),
}
@@ -235,6 +268,7 @@ def train_one_epoch(
# ---------------------------------------------------------------------------
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)
@@ -248,7 +282,7 @@ def main(args):
with open(output_dir / "config.json", "w") as f:
json.dump(config, f, indent=2, default=str)
train_loader, test_loader = build_dataloaders(
train_loader, _ = build_dataloaders(
data_root=args.data_root,
descriptions_path=args.descriptions_path,
text_levels=args.text_levels,
@@ -257,6 +291,18 @@ def main(args):
batch_size=args.micro_batch_size,
num_workers=args.num_workers,
image_size=args.image_size,
build_test=False, # оценка идёт через multi-positive eval ниже
)
# 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(
@@ -292,11 +338,26 @@ def main(args):
betas=(0.9, 0.98),
)
scheduler = CosineAnnealingLR(
# 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,
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)
@@ -341,21 +402,25 @@ def main(args):
scheduler.step()
if epoch % args.eval_every == 0 or epoch == args.epochs:
eval_metrics = evaluate(model, test_loader, device, amp_ctx)
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 = {}
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 | AP=%.3f",
"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("AP", 0),
eval_metrics.get("mAP", 0),
)
if epoch == 1:
@@ -382,6 +447,11 @@ def main(args):
},
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(
@@ -434,7 +504,7 @@ def parse_args():
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.0)
p.add_argument("--label_smoothing", type=float, default=0.1)
p.add_argument("--eval_every", type=int, default=1)
# Performance
@@ -447,9 +517,12 @@ def parse_args():
# 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())
main(parse_args())