da2d2ea90e
- Single DINOv3 WEB for both drone and satellite branches (shared_encoder=True default) - One set of MONA adapters instead of two: 7M trainable vs 14M - Total params: 438M (was 748M), trainable: 10.6M (was 17.6M) - Asymmetric mode still available via shared_encoder=False - Add gradient accumulation (grad_accum_steps, --grad-accum CLI flag) - Update model summary in README Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
835 lines
29 KiB
Python
835 lines
29 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).
|
|
|
|
Supports gin-config, W&B, TensorBoard, Grad-CAM, gradient monitoring,
|
|
PyTorch Profiler, and torchinfo model summary.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import math
|
|
import time
|
|
import warnings
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
import coloredlogs
|
|
import gin
|
|
import pandas as pd
|
|
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.training.plot_metrics import generate_plots
|
|
from src.training.trackers import ExperimentTracker
|
|
from src.training.grad_monitor import compute_gradient_norms, log_gradient_summary
|
|
from src.training.profiling import TrainingProfiler, print_model_summary
|
|
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"
|
|
|
|
|
|
@gin.configurable(module="src.training.train_gtauav")
|
|
@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
|
|
shared_encoder: bool = True # single DINOv3 WEB for both branches (saves ~4-5 GB)
|
|
|
|
# Training.
|
|
resume_from: str | None = None # path to checkpoint for resuming
|
|
output_dir: str = "out/gtauav/with_text"
|
|
epochs: int = 10
|
|
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
|
|
grad_accum_steps: int = 1 # gradient accumulation steps (effective_batch = batch_size * accum)
|
|
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
|
|
|
|
# Tracking & diagnostics.
|
|
use_wandb: bool = False
|
|
use_tb: bool = True
|
|
wandb_project: str = "caption-test-gtauav"
|
|
wandb_run_name: str | None = None
|
|
wandb_entity: str | None = None
|
|
log_grad_norms: bool = True
|
|
use_gradcam: bool = False
|
|
gradcam_every: int = 5 # Grad-CAM every N epochs
|
|
gradcam_samples: int = 8
|
|
use_profiler: bool = False
|
|
profiler_warmup: int = 3
|
|
profiler_active: int = 5
|
|
|
|
|
|
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=" eval", 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
|
|
|
|
|
|
class CSVLogger:
|
|
"""Log train/val metrics to CSV files using pandas.
|
|
|
|
Creates:
|
|
{output_dir}/logs/train.csv — epoch-level train averages
|
|
{output_dir}/logs/val.csv — epoch-level val metrics
|
|
{output_dir}/logs/train_batches.csv — per-batch train metrics (all epochs)
|
|
{output_dir}/logs/epoch_{N}_train.csv — per-epoch summary
|
|
{output_dir}/logs/epoch_{N}_val.csv — per-epoch val
|
|
{output_dir}/logs/epoch_{N}_batches.csv — per-batch for single epoch
|
|
"""
|
|
|
|
def __init__(self, output_dir: Path) -> None:
|
|
self.log_dir = output_dir / "logs"
|
|
self.log_dir.mkdir(parents=True, exist_ok=True)
|
|
self.train_rows: list[dict] = []
|
|
self.val_rows: list[dict] = []
|
|
self._current_epoch: int = -1
|
|
self._batch_columns: list[str] | None = None
|
|
self._cumulative_batch_path = self.log_dir / "train_batches.csv"
|
|
self._epoch_batch_path: Path | None = None
|
|
|
|
def log_batch(self, epoch: int, batch_idx: int, global_step: int, metrics: dict) -> None:
|
|
"""Log metrics for a single training batch. Writes to disk immediately."""
|
|
row = {"epoch": epoch, "batch": batch_idx, "global_step": global_step, **metrics}
|
|
|
|
# On new epoch, start a fresh per-epoch CSV.
|
|
if epoch != self._current_epoch:
|
|
self._current_epoch = epoch
|
|
self._epoch_batch_path = self.log_dir / f"epoch_{epoch:03d}_batches.csv"
|
|
|
|
# Determine columns on first call (consistent order).
|
|
if self._batch_columns is None:
|
|
self._batch_columns = list(row.keys())
|
|
|
|
row_df = pd.DataFrame([row], columns=self._batch_columns)
|
|
write_header = not self._cumulative_batch_path.exists()
|
|
|
|
# Append to cumulative CSV.
|
|
row_df.to_csv(
|
|
self._cumulative_batch_path, mode="a", header=write_header, index=False,
|
|
)
|
|
# Append to per-epoch CSV.
|
|
write_epoch_header = not self._epoch_batch_path.exists()
|
|
row_df.to_csv(
|
|
self._epoch_batch_path, mode="a", header=write_epoch_header, index=False,
|
|
)
|
|
|
|
def log_train(self, epoch: int, metrics: dict, lr: float, elapsed: float) -> None:
|
|
"""Log epoch-level train averages."""
|
|
row = {"epoch": epoch, "lr": lr, "elapsed_s": round(elapsed, 1), **metrics}
|
|
self.train_rows.append(row)
|
|
pd.DataFrame(self.train_rows).to_csv(self.log_dir / "train.csv", index=False)
|
|
pd.DataFrame([row]).to_csv(self.log_dir / f"epoch_{epoch:03d}_train.csv", index=False)
|
|
|
|
def log_val(self, epoch: int, metrics: dict) -> None:
|
|
row = {"epoch": epoch, **metrics}
|
|
self.val_rows.append(row)
|
|
pd.DataFrame(self.val_rows).to_csv(self.log_dir / "val.csv", index=False)
|
|
pd.DataFrame([row]).to_csv(self.log_dir / f"epoch_{epoch:03d}_val.csv", index=False)
|
|
|
|
|
|
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)
|
|
|
|
# --- Experiment tracker (W&B + TensorBoard) ---
|
|
tracker = ExperimentTracker(
|
|
output_dir=output_dir,
|
|
config=vars(cfg),
|
|
use_wandb=cfg.use_wandb,
|
|
use_tb=cfg.use_tb,
|
|
wandb_project=cfg.wandb_project,
|
|
wandb_run_name=cfg.wandb_run_name,
|
|
wandb_entity=cfg.wandb_entity,
|
|
)
|
|
|
|
# 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)"
|
|
enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)"
|
|
LOGGER.info("Building model — %s, %s", mode_str, enc_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,
|
|
shared_encoder=cfg.shared_encoder,
|
|
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:,}",
|
|
)
|
|
|
|
# --- Model summary (torchinfo) ---
|
|
model_summary = print_model_summary(model, device=cfg.device)
|
|
(output_dir / "model_summary.txt").write_text(model_summary)
|
|
|
|
# --- W&B model watching (gradient + weight histograms) ---
|
|
if tracker.has_wandb:
|
|
tracker.watch_model(model, log_freq=50)
|
|
|
|
# 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,
|
|
)
|
|
|
|
effective_batch = cfg.batch_size * cfg.grad_accum_steps
|
|
LOGGER.info(
|
|
"train=%d test=%d batch=%d accum=%d effective_batch=%d",
|
|
len(train_ds), len(test_ds), cfg.batch_size, cfg.grad_accum_steps, effective_batch,
|
|
)
|
|
|
|
# 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 (counted in optimizer steps).
|
|
steps_per_epoch = math.ceil(len(train_loader) / cfg.grad_accum_steps)
|
|
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] = []
|
|
csv_logger = CSVLogger(output_dir)
|
|
|
|
# --- Optional profiler (first epoch only) ---
|
|
profiler = None
|
|
if cfg.use_profiler and start_epoch == 0:
|
|
profiler = TrainingProfiler(
|
|
output_dir=output_dir,
|
|
n_warmup=cfg.profiler_warmup,
|
|
n_active=cfg.profiler_active,
|
|
)
|
|
profiler.start()
|
|
|
|
LOGGER.info("Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch)
|
|
|
|
global_step = start_epoch * steps_per_epoch
|
|
best_r1 = 0.0
|
|
|
|
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,
|
|
)
|
|
accum = cfg.grad_accum_steps
|
|
for batch in pbar:
|
|
# Zero gradients only at the start of each accumulation window.
|
|
if n_batches % accum == 0:
|
|
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,
|
|
)
|
|
|
|
# Scale loss by accumulation steps so gradients average correctly.
|
|
total_loss = loss_dict["total"] / accum
|
|
scaler.scale(total_loss).backward()
|
|
|
|
# Optimizer step only after accumulating `accum` micro-batches.
|
|
is_accum_step = (n_batches + 1) % accum == 0 or (n_batches + 1) == len(train_loader)
|
|
if is_accum_step:
|
|
if cfg.grad_clip > 0:
|
|
scaler.unscale_(optimizer)
|
|
nn.utils.clip_grad_norm_(
|
|
model.trainable_parameters(),
|
|
max_norm=cfg.grad_clip,
|
|
)
|
|
|
|
# --- Gradient monitoring (after unscale, before step) ---
|
|
if cfg.log_grad_norms and n_batches % (50 * accum) < accum:
|
|
grad_norms = compute_gradient_norms(model, loss_fn)
|
|
tracker.log_gradients(epoch, grad_norms, step=global_step)
|
|
if n_batches < accum:
|
|
log_gradient_summary(grad_norms)
|
|
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
|
|
scheduler.step()
|
|
global_step += 1
|
|
|
|
# --- Per-batch tracking (log unscaled loss) ---
|
|
raw_loss = total_loss.item() * accum # undo /accum for logging
|
|
step_metrics = {
|
|
"loss": raw_loss,
|
|
"temperature": float(loss_dict["temperature"].item()),
|
|
"gate_q": float(loss_dict["gate_q"].item()),
|
|
"gate_g": float(loss_dict["gate_g"].item()),
|
|
"lr": optimizer.param_groups[0]["lr"],
|
|
}
|
|
tracker.log_train(epoch, step_metrics, step=global_step)
|
|
csv_logger.log_batch(epoch, n_batches, global_step, step_metrics)
|
|
|
|
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"{raw_loss:.3f}",
|
|
tau=f"{loss_dict['temperature'].item():.4f}",
|
|
gq=f"{loss_dict['gate_q'].item():.3f}",
|
|
gg=f"{loss_dict['gate_g'].item():.3f}",
|
|
)
|
|
|
|
# --- Profiler step ---
|
|
if profiler is not None:
|
|
profiler.step()
|
|
if profiler.is_done(n_batches):
|
|
profiler.export()
|
|
profiler = None
|
|
|
|
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,
|
|
}
|
|
|
|
# Log train metrics to CSV + generate plots every epoch.
|
|
csv_logger.log_train(epoch, means, optimizer.param_groups[0]["lr"], elapsed)
|
|
generate_plots(csv_logger.log_dir)
|
|
|
|
# --- Log VRAM usage ---
|
|
if torch.cuda.is_available():
|
|
vram_gb = torch.cuda.max_memory_allocated() / 1e9
|
|
tracker.log_scalar("system/vram_peak_gb", vram_gb, step=global_step)
|
|
|
|
# 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
|
|
csv_logger.log_val(epoch, val_metrics)
|
|
generate_plots(csv_logger.log_dir)
|
|
tracker.log_val(epoch, val_metrics, step=global_step)
|
|
|
|
# Track best R@1.
|
|
r1 = val_metrics.get("r@1_q2g", 0.0)
|
|
if r1 > best_r1:
|
|
best_r1 = r1
|
|
tracker.log_scalar("val/best_r@1_q2g", best_r1, step=global_step)
|
|
|
|
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),
|
|
)
|
|
|
|
# --- Grad-CAM visualization ---
|
|
if cfg.use_gradcam and (epoch + 1) % cfg.gradcam_every == 0:
|
|
from src.training.gradcam import generate_gradcam_samples
|
|
overlays = generate_gradcam_samples(
|
|
model=model,
|
|
dataloader=test_loader,
|
|
device=cfg.device,
|
|
output_dir=str(output_dir),
|
|
n_samples=cfg.gradcam_samples,
|
|
epoch=epoch,
|
|
)
|
|
# Log first few overlays to tracker.
|
|
for i, overlay in enumerate(overlays[:4]):
|
|
kind = "drone" if i % 2 == 0 else "sat"
|
|
tracker.log_image(
|
|
f"gradcam/{kind}_{i//2}",
|
|
overlay,
|
|
step=global_step,
|
|
caption=f"Epoch {epoch} {kind} Grad-CAM",
|
|
)
|
|
|
|
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)
|
|
|
|
# --- Log final summary to W&B ---
|
|
tracker.log_summary({
|
|
"best_r@1_q2g": best_r1,
|
|
"final_r@1_q2g": final_metrics.get("r@1_q2g", 0.0),
|
|
"final_r@5_q2g": final_metrics.get("r@5_q2g", 0.0),
|
|
"final_r@10_q2g": final_metrics.get("r@10_q2g", 0.0),
|
|
"final_gate_q": final_metrics.get("gate_q", 1.0),
|
|
"final_gate_g": final_metrics.get("gate_g", 1.0),
|
|
})
|
|
|
|
# --- Cleanup profiler if still running ---
|
|
if profiler is not None:
|
|
profiler.export()
|
|
|
|
tracker.close()
|
|
|
|
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(
|
|
"--config", type=str, default=None,
|
|
help="Path to gin config file (e.g. conf/gtauav_balanced.gin).",
|
|
)
|
|
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=None,
|
|
help="Batch size.",
|
|
)
|
|
parser.add_argument(
|
|
"--grad-accum", type=int, default=None,
|
|
help="Gradient accumulation steps (effective_batch = batch_size * accum).",
|
|
)
|
|
parser.add_argument(
|
|
"--epochs", type=int, default=None,
|
|
help="Number of epochs.",
|
|
)
|
|
parser.add_argument(
|
|
"--lr", type=float, default=None,
|
|
help="Learning rate for projections.",
|
|
)
|
|
parser.add_argument(
|
|
"--text-lr-factor", type=float, default=None,
|
|
help="Text encoder LR = lr * factor (default 0.1 = 10x lower).",
|
|
)
|
|
parser.add_argument(
|
|
"--warmup-epochs", type=int, default=None,
|
|
help="Linear warmup epochs.",
|
|
)
|
|
parser.add_argument(
|
|
"--init-gate", type=float, default=None,
|
|
help="Initial gate value (image weight).",
|
|
)
|
|
# Tracking flags.
|
|
parser.add_argument("--wandb", action="store_true", help="Enable W&B tracking.")
|
|
parser.add_argument("--no-tb", action="store_true", help="Disable TensorBoard.")
|
|
parser.add_argument("--gradcam", action="store_true", help="Enable Grad-CAM visualization.")
|
|
parser.add_argument("--profile", action="store_true", help="Enable PyTorch profiler (first epoch).")
|
|
parser.add_argument("--no-grad-norms", action="store_true", help="Disable gradient norm logging.")
|
|
# Gin overrides.
|
|
parser.add_argument(
|
|
"--gin-param", type=str, nargs="*", default=[],
|
|
help="Gin parameter overrides (e.g. 'TrainConfigGTAUAV.epochs=30').",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Parse gin config if provided.
|
|
if args.config is not None:
|
|
gin.parse_config_file(args.config)
|
|
if args.gin_param:
|
|
gin.parse_config(args.gin_param)
|
|
|
|
# Create config (gin bindings apply via @gin.configurable).
|
|
cfg = TrainConfigGTAUAV()
|
|
|
|
# CLI overrides take priority over gin.
|
|
if args.baseline:
|
|
cfg.baseline_mode = True
|
|
if args.resume is not None:
|
|
cfg.resume_from = args.resume
|
|
if args.batch_size is not None:
|
|
cfg.batch_size = args.batch_size
|
|
if args.grad_accum is not None:
|
|
cfg.grad_accum_steps = args.grad_accum
|
|
if args.epochs is not None:
|
|
cfg.epochs = args.epochs
|
|
if args.lr is not None:
|
|
cfg.learning_rate = args.lr
|
|
if args.text_lr_factor is not None:
|
|
cfg.text_lr_factor = args.text_lr_factor
|
|
if args.warmup_epochs is not None:
|
|
cfg.warmup_epochs = args.warmup_epochs
|
|
if args.init_gate is not None:
|
|
cfg.init_gate = args.init_gate
|
|
if args.filter_meta is not None:
|
|
cfg.filter_meta = args.filter_meta
|
|
|
|
# Tracking overrides.
|
|
if args.wandb:
|
|
cfg.use_wandb = True
|
|
if args.no_tb:
|
|
cfg.use_tb = False
|
|
if args.gradcam:
|
|
cfg.use_gradcam = True
|
|
if args.profile:
|
|
cfg.use_profiler = True
|
|
if args.no_grad_norms:
|
|
cfg.log_grad_norms = False
|
|
|
|
if args.output_dir is not None:
|
|
cfg.output_dir = args.output_dir
|
|
elif args.baseline and args.output_dir is None:
|
|
cfg.output_dir = "out/gtauav/baseline"
|
|
|
|
train(cfg)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|