455 lines
16 KiB
Python
455 lines
16 KiB
Python
"""Цикл обучения 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 (веса не общие, т.к. визуальные
|
||
домены различаются).
|
||
|
||
Оптимизировано под 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 \\
|
||
--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 sys
|
||
import time
|
||
from contextlib import nullcontext
|
||
from pathlib import Path
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
from torch.optim import AdamW
|
||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||
|
||
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.data.gta_uav import build_dataloaders
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s %(name)s %(levelname)s %(message)s",
|
||
)
|
||
LOGGER = logging.getLogger("cvgl.train")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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()
|
||
def evaluate(
|
||
model: nn.Module,
|
||
test_loader,
|
||
device: torch.device,
|
||
amp_ctx,
|
||
) -> dict[str, float]:
|
||
"""Прогнать test set: drone (fused) vs satellite (fused) — симметрично."""
|
||
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],
|
||
)
|
||
|
||
model.train()
|
||
return metrics
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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_t2i = 0.0
|
||
total_acc_i2t = 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_t2i += loss_dict["acc_t2i"].item()
|
||
total_acc_i2t += loss_dict["acc_i2t"].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_t2i"].item(),
|
||
loss_dict["acc_i2t"].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_t2i": total_acc_t2i / max(n_steps, 1),
|
||
"acc_i2t": total_acc_i2t / max(n_steps, 1),
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def main(args):
|
||
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, test_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,
|
||
)
|
||
|
||
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),
|
||
)
|
||
|
||
scheduler = CosineAnnealingLR(
|
||
optimizer,
|
||
T_max=args.epochs,
|
||
eta_min=args.lr * 0.01,
|
||
)
|
||
|
||
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 = evaluate(model, test_loader, device, amp_ctx)
|
||
else:
|
||
eval_metrics = {}
|
||
|
||
elapsed = time.time() - t0
|
||
|
||
LOGGER.info(
|
||
"📈 Epoch %d/%d (%.0fs) | loss=%.4f | "
|
||
"R@1=%.3f R@5=%.3f R@10=%.3f | AP=%.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),
|
||
)
|
||
|
||
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",
|
||
)
|
||
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=384)
|
||
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)
|
||
|
||
# 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.0)
|
||
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")
|
||
|
||
return p.parse_args()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main(parse_args()) |