initial commit
This commit is contained in:
pikaliov
455
train.py
Normal file
455
train.py
Normal file
@@ -0,0 +1,455 @@
|
||||
"""Цикл обучения 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())
|
||||
Reference in New Issue
Block a user