1297 lines
49 KiB
Python
1297 lines
49 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
|
||
import torch.nn.functional as F
|
||
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,
|
||
GTAUAVDroneQuery,
|
||
GTAUAVSatGallery,
|
||
collate_drone_query,
|
||
collate_gtauav_batch,
|
||
collate_sat_gallery,
|
||
)
|
||
from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler
|
||
from src.datasets.embedding_cache import EmbeddingCache
|
||
from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler
|
||
from src.losses.multi_infonce import InfoNCELoss
|
||
from src.losses.weighted_infonce import WeightedInfoNCELoss
|
||
from src.losses.hard_negatives import NegativeMemoryBank
|
||
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 (simpler, half the params)
|
||
mona_bottleneck: int = 64
|
||
mona_last_n_blocks: int = 12 # inject adapters only in last 12 of 24 ViT blocks
|
||
gradient_checkpointing: bool = True # trade compute for VRAM (allows larger batch)
|
||
# StripNet backbone option (replaces DINOv3 when backbone="stripnet").
|
||
backbone: str = "dinov3" # "dinov3" or "stripnet"
|
||
stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth"
|
||
stripnet_mona_last_n_stages: int = 2 # Conv-MONA in last N of 4 StripNet stages (0 = disable MONA)
|
||
stripnet_freeze: bool = True # If False, StripNet backbone is fully trainable (full fine-tune)
|
||
stripnet_backbone_lr_factor: float = 0.1 # Backbone LR = learning_rate * factor (only when unfrozen)
|
||
|
||
# Training.
|
||
resume_from: str | None = None # path to checkpoint for resuming
|
||
# output_dir: str = "out/gtauav/with_text"
|
||
output_dir: str = "out/gtauav/with_text_exp_gate_SRGF"
|
||
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.
|
||
loss_type: str = "symmetric" # "symmetric" (InfoNCE) or "weighted" (WeightedInfoNCE)
|
||
tau_init: float = 0.07
|
||
label_smoothing: float = 0.1
|
||
learnable_temperature: bool = True
|
||
weight_q2g: float = 0.6
|
||
weight_g2q: float = 0.4
|
||
neg_bank_size: int = 4096 # hard negative memory bank size (0 = disabled)
|
||
|
||
# Sampling.
|
||
sampler_type: str = "mutex" # "mutex" (no false negatives) or "dss" (DSS + mutex)
|
||
dss_reembed_every: int = 1 # Re-embed train queries every N epochs for DSS.
|
||
dss_warmup_epochs: int = 1 # Use mutex-only for the first N epochs (fresh model embeddings aren't useful)
|
||
dss_knn_device: str = "cuda" # Device for similarity matmul in DSS sampler.
|
||
dss_use_lsh: bool = False # Approximate kNN via LSH (opt-in; exact is fast at 25K).
|
||
dss_lsh_num_tables: int = 8
|
||
dss_lsh_num_bits: int = 14
|
||
dss_cache_dir: str | None = None # Disk cache for embeddings; None = disabled.
|
||
# Legacy alias kept for backward compatibility.
|
||
use_mutex_sampler: 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,
|
||
stripnet_backbone_lr_factor: float = 0.1,
|
||
) -> list[dict]:
|
||
"""Build optimizer param groups with separate LR for text encoder and unfrozen StripNet backbone.
|
||
|
||
Groups:
|
||
- text_encoder.* → lr * text_lr_factor (default 1e-5)
|
||
- image_encoder.backbone.* (when StripNet unfrozen) → lr * stripnet_backbone_lr_factor (default 1e-5)
|
||
- everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv) → lr
|
||
"""
|
||
text_params = []
|
||
backbone_params = []
|
||
other_params = []
|
||
|
||
is_stripnet = isinstance(getattr(model, "image_encoder", None), nn.Module) and \
|
||
getattr(model, "backbone", "dinov3") == "stripnet"
|
||
|
||
for name, param in model.named_parameters():
|
||
if not param.requires_grad:
|
||
continue
|
||
if "text_encoder" in name:
|
||
text_params.append(param)
|
||
elif is_stripnet and name.startswith("image_encoder.backbone.") and "mona_" not in name:
|
||
backbone_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})
|
||
if backbone_params:
|
||
groups.append({"params": backbone_params, "lr": lr * stripnet_backbone_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 _embed_drone_queries(
|
||
model: AsymmetricEncoder,
|
||
train_ds: GTAUAVDataset,
|
||
device: str,
|
||
batch_size: int,
|
||
num_workers: int,
|
||
) -> torch.Tensor:
|
||
"""Forward all drone queries and return [N, D] embeddings on CPU.
|
||
|
||
Used by DynamicSimilaritySampler to rank drones by visual similarity.
|
||
Runs with model.eval() but restores original train state afterwards.
|
||
"""
|
||
was_training = model.training
|
||
model.eval()
|
||
|
||
query_ds = GTAUAVDroneQuery(train_ds)
|
||
loader = DataLoader(
|
||
query_ds,
|
||
batch_size=batch_size,
|
||
shuffle=False,
|
||
num_workers=num_workers,
|
||
collate_fn=collate_drone_query,
|
||
pin_memory=True,
|
||
)
|
||
|
||
embs: list[torch.Tensor] = []
|
||
for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False):
|
||
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
||
q = model.encode_query(
|
||
drone_img,
|
||
batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
|
||
)
|
||
embs.append(q.cpu())
|
||
|
||
if was_training:
|
||
model.train()
|
||
return torch.cat(embs, dim=0)
|
||
|
||
|
||
@torch.no_grad()
|
||
def _evaluate(
|
||
model: AsymmetricEncoder,
|
||
loader: DataLoader,
|
||
device: str,
|
||
loss_fn: nn.Module | None = None,
|
||
epoch: int = 0,
|
||
total_epochs: int = 1,
|
||
k_values: tuple[int, ...] = (1, 5, 10),
|
||
max_batches: int | None = None,
|
||
desc: str = "eval",
|
||
) -> dict[str, float]:
|
||
"""Compute R@K and MRR on the full satellite gallery.
|
||
|
||
Standard CVGL retrieval: forward every unique satellite in the dataset
|
||
once (gallery), forward every drone query, then rank gallery by
|
||
cosine similarity. A query counts as a hit@K if ANY of its valid
|
||
satellite matches (pair_pos_sate_img_list ∪ pair_pos_semipos_sate_img_list)
|
||
appears in the top-K.
|
||
|
||
`max_batches` subsamples the drone queries (not the gallery) — useful
|
||
for a quick train-side sanity check.
|
||
"""
|
||
dataset = loader.dataset
|
||
if not isinstance(dataset, GTAUAVDataset):
|
||
raise TypeError(f"_evaluate expects GTAUAVDataset, got {type(dataset).__name__}")
|
||
|
||
model.eval()
|
||
|
||
batch_size = loader.batch_size or 32
|
||
num_workers = getattr(loader, "num_workers", 0)
|
||
pin_memory = getattr(loader, "pin_memory", False)
|
||
|
||
gallery_ds = GTAUAVSatGallery(dataset)
|
||
query_ds = GTAUAVDroneQuery(dataset)
|
||
|
||
gallery_loader = DataLoader(
|
||
gallery_ds,
|
||
batch_size=batch_size,
|
||
shuffle=False,
|
||
num_workers=num_workers,
|
||
pin_memory=pin_memory,
|
||
collate_fn=collate_sat_gallery,
|
||
)
|
||
query_loader = DataLoader(
|
||
query_ds,
|
||
batch_size=batch_size,
|
||
shuffle=False,
|
||
num_workers=num_workers,
|
||
pin_memory=pin_memory,
|
||
collate_fn=collate_drone_query,
|
||
)
|
||
|
||
# --- Gallery forward (all unique sats) ---
|
||
gallery_embs: list[torch.Tensor] = []
|
||
gallery_names: list[str] = []
|
||
for batch in tqdm(gallery_loader, desc=f" {desc}-gallery", unit="batch", leave=False):
|
||
sat_img = batch["sat_img"].to(device, non_blocking=True)
|
||
g = model.encode_gallery(
|
||
sat_img,
|
||
batch["sat_caption_l1"], batch["sat_caption_l2"], batch["sat_caption_l3"],
|
||
)
|
||
gallery_embs.append(g.cpu())
|
||
gallery_names.extend(batch["sat_names"])
|
||
gallery = torch.cat(gallery_embs, dim=0) # [N_sat, D]
|
||
|
||
# --- Query forward (optionally subsampled via max_batches) ---
|
||
query_embs: list[torch.Tensor] = []
|
||
query_valid_names: list[list[str]] = []
|
||
batch_losses: list[float] = []
|
||
sat_name_to_idx: dict[str, int] = {name: i for i, name in enumerate(gallery_names)}
|
||
|
||
for i, batch in enumerate(tqdm(query_loader, desc=f" {desc}-query", unit="batch", leave=False)):
|
||
if max_batches is not None and i >= max_batches:
|
||
break
|
||
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
||
q = model.encode_query(
|
||
drone_img,
|
||
batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
|
||
)
|
||
query_embs.append(q.cpu())
|
||
query_valid_names.extend(batch["valid_sat_names"])
|
||
|
||
# Per-batch loss: use first valid sat per query as its paired gallery.
|
||
if loss_fn is not None:
|
||
pair_indices: list[int] = []
|
||
for names in batch["valid_sat_names"]:
|
||
for name in names:
|
||
if name in sat_name_to_idx:
|
||
pair_indices.append(sat_name_to_idx[name])
|
||
break
|
||
else:
|
||
pair_indices.append(-1)
|
||
if all(idx >= 0 for idx in pair_indices):
|
||
paired_gallery = gallery[pair_indices].to(device)
|
||
fake_embeddings = {
|
||
"query": q,
|
||
"gallery": paired_gallery,
|
||
"gate_q": model.fusion_query.gate_value,
|
||
"gate_g": model.fusion_gallery.gate_value,
|
||
}
|
||
loss_dict = loss_fn(fake_embeddings, epoch=epoch, total_epochs=total_epochs)
|
||
batch_losses.append(float(loss_dict["total"].item()))
|
||
|
||
query = torch.cat(query_embs, dim=0) # [N_q, D]
|
||
n_query = query.size(0)
|
||
|
||
# --- Similarity + rankings ---
|
||
sim = query @ gallery.t() # [N_q, N_sat]
|
||
sorted_idx = sim.argsort(dim=1, descending=True)
|
||
|
||
metrics: dict[str, float] = {}
|
||
if batch_losses:
|
||
metrics["loss"] = sum(batch_losses) / len(batch_losses)
|
||
|
||
# Precompute valid gallery index sets per query.
|
||
valid_idx_per_query: list[set[int]] = []
|
||
for names in query_valid_names:
|
||
valid = {sat_name_to_idx[n] for n in names if n in sat_name_to_idx}
|
||
valid_idx_per_query.append(valid)
|
||
|
||
# R@K with multi-match.
|
||
for k in k_values:
|
||
hits = 0
|
||
for i in range(n_query):
|
||
top_k = set(sorted_idx[i, :k].tolist())
|
||
if valid_idx_per_query[i] & top_k:
|
||
hits += 1
|
||
metrics[f"r@{k}_q2g"] = hits / max(n_query, 1)
|
||
|
||
# MRR over valid matches (kept key `ap_q2g` for CSV/plot compatibility).
|
||
mrr_sum = 0.0
|
||
n_scored = 0
|
||
for i in range(n_query):
|
||
valid = valid_idx_per_query[i]
|
||
if not valid:
|
||
continue
|
||
n_scored += 1
|
||
for rank, gidx in enumerate(sorted_idx[i].tolist()):
|
||
if gidx in valid:
|
||
mrr_sum += 1.0 / (rank + 1)
|
||
break
|
||
metrics["ap_q2g"] = mrr_sum / max(n_scored, 1)
|
||
|
||
# --- g2q (satellite → drone): invert ground-truth ---
|
||
n_gallery = gallery.size(0)
|
||
valid_q_per_sat: list[set[int]] = [set() for _ in range(n_gallery)]
|
||
for q_idx, gset in enumerate(valid_idx_per_query):
|
||
for g_idx in gset:
|
||
valid_q_per_sat[g_idx].add(q_idx)
|
||
|
||
sorted_idx_g2q = sim.t().argsort(dim=1, descending=True) # [N_sat, n_query]
|
||
n_scored_g2q = sum(1 for s in valid_q_per_sat if s)
|
||
|
||
for k in k_values:
|
||
hits_g2q = 0
|
||
for i in range(n_gallery):
|
||
valid = valid_q_per_sat[i]
|
||
if not valid:
|
||
continue
|
||
top_k = set(sorted_idx_g2q[i, :k].tolist())
|
||
if valid & top_k:
|
||
hits_g2q += 1
|
||
metrics[f"r@{k}_g2q"] = hits_g2q / max(n_scored_g2q, 1)
|
||
|
||
mrr_sum_g2q = 0.0
|
||
for i in range(n_gallery):
|
||
valid = valid_q_per_sat[i]
|
||
if not valid:
|
||
continue
|
||
for rank, qidx in enumerate(sorted_idx_g2q[i].tolist()):
|
||
if qidx in valid:
|
||
mrr_sum_g2q += 1.0 / (rank + 1)
|
||
break
|
||
metrics["ap_g2q"] = mrr_sum_g2q / max(n_scored_g2q, 1)
|
||
|
||
metrics["n_query"] = float(n_query)
|
||
metrics["n_gallery"] = float(n_gallery)
|
||
metrics["n_scored_g2q"] = float(n_scored_g2q)
|
||
|
||
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}_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._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
|
||
|
||
# Load existing CSV data on resume (so plots show full history).
|
||
train_csv = self.log_dir / "train.csv"
|
||
val_csv = self.log_dir / "val.csv"
|
||
train_recall_csv = self.log_dir / "train_recall.csv"
|
||
if train_csv.exists():
|
||
self.train_rows = pd.read_csv(train_csv).to_dict("records")
|
||
LOGGER.info("CSVLogger: loaded %d previous train epochs", len(self.train_rows))
|
||
else:
|
||
self.train_rows = []
|
||
if val_csv.exists():
|
||
self.val_rows = pd.read_csv(val_csv).to_dict("records")
|
||
LOGGER.info("CSVLogger: loaded %d previous val epochs", len(self.val_rows))
|
||
else:
|
||
self.val_rows = []
|
||
if train_recall_csv.exists():
|
||
self.train_recall_rows = pd.read_csv(train_recall_csv).to_dict("records")
|
||
else:
|
||
self.train_recall_rows = []
|
||
|
||
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. Replaces existing entry for same epoch on resume."""
|
||
row = {"epoch": epoch, "lr": lr, "elapsed_s": round(elapsed, 1), **metrics}
|
||
# Remove previous entry for this epoch (resume may re-run it).
|
||
self.train_rows = [r for r in self.train_rows if r.get("epoch") != epoch]
|
||
self.train_rows.append(row)
|
||
pd.DataFrame(self.train_rows).to_csv(self.log_dir / "train.csv", index=False)
|
||
|
||
def log_val(self, epoch: int, metrics: dict) -> None:
|
||
"""Log val metrics. Replaces existing entry for same epoch on resume."""
|
||
row = {"epoch": epoch, **metrics}
|
||
self.val_rows = [r for r in self.val_rows if r.get("epoch") != epoch]
|
||
self.val_rows.append(row)
|
||
pd.DataFrame(self.val_rows).to_csv(self.log_dir / "val.csv", index=False)
|
||
|
||
def log_train_recall(self, epoch: int, metrics: dict) -> None:
|
||
"""Log train recall metrics. Replaces existing entry for same epoch."""
|
||
row = {"epoch": epoch, **metrics}
|
||
self.train_recall_rows = [r for r in self.train_recall_rows if r.get("epoch") != epoch]
|
||
self.train_recall_rows.append(row)
|
||
pd.DataFrame(self.train_recall_rows).to_csv(self.log_dir / "train_recall.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)"
|
||
if cfg.backbone == "stripnet":
|
||
enc_str = "StripNet-small (shared, 512→1024 proj)"
|
||
else:
|
||
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,
|
||
mona_bottleneck=cfg.mona_bottleneck,
|
||
mona_last_n_blocks=cfg.mona_last_n_blocks,
|
||
device=cfg.device,
|
||
backbone=cfg.backbone,
|
||
stripnet_path=cfg.stripnet_path,
|
||
stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages,
|
||
stripnet_freeze=cfg.stripnet_freeze,
|
||
).to(cfg.device)
|
||
LOGGER.info("embed_dim=%d", model.embed_dim)
|
||
|
||
# --- Gradient checkpointing (trade compute for VRAM) ---
|
||
# StripNet doesn't expose set_gradient_checkpointing — skip silently.
|
||
if cfg.gradient_checkpointing and cfg.backbone == "dinov3":
|
||
if cfg.shared_encoder:
|
||
model.image_encoder.set_gradient_checkpointing(True)
|
||
else:
|
||
model.drone_encoder.set_gradient_checkpointing(True)
|
||
model.sat_encoder.set_gradient_checkpointing(True)
|
||
if model.text_encoder is not None:
|
||
model.text_encoder.transformer.gradient_checkpointing = True
|
||
LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)")
|
||
elif cfg.gradient_checkpointing and cfg.backbone == "stripnet":
|
||
if model.text_encoder is not None:
|
||
model.text_encoder.transformer.gradient_checkpointing = True
|
||
LOGGER.info("Gradient checkpointing enabled (DGTRS only; StripNet doesn't support)")
|
||
|
||
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.
|
||
if cfg.loss_type == "symmetric":
|
||
loss_fn = InfoNCELoss(
|
||
temperature_init=cfg.tau_init,
|
||
learnable_temperature=cfg.learnable_temperature,
|
||
label_smoothing=cfg.label_smoothing,
|
||
weight_q2g=cfg.weight_q2g,
|
||
weight_g2q=cfg.weight_g2q,
|
||
)
|
||
loss_name = "SymmetricInfoNCE"
|
||
elif cfg.loss_type == "weighted":
|
||
loss_fn = WeightedInfoNCELoss(
|
||
temperature_init=cfg.tau_init,
|
||
learnable_temperature=cfg.learnable_temperature,
|
||
label_smoothing=cfg.label_smoothing,
|
||
)
|
||
loss_name = "WeightedInfoNCE"
|
||
else:
|
||
raise ValueError(f"Unknown loss_type={cfg.loss_type!r} (expected 'symmetric' or 'weighted')")
|
||
|
||
LOGGER.info(
|
||
"Loss: %s Temperature: %s (init=%.3f) q2g=%.2f g2q=%.2f",
|
||
loss_name,
|
||
"learnable" if cfg.learnable_temperature else "fixed",
|
||
cfg.tau_init, cfg.weight_q2g, cfg.weight_g2q,
|
||
)
|
||
|
||
# Hard negative memory bank.
|
||
neg_bank = None
|
||
if cfg.neg_bank_size > 0:
|
||
neg_bank = NegativeMemoryBank(size=cfg.neg_bank_size, dim=model.embed_dim).to(cfg.device)
|
||
LOGGER.info("Negative memory bank: size=%d, dim=%d", cfg.neg_bank_size, model.embed_dim)
|
||
|
||
# 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,
|
||
)
|
||
|
||
sat_cand_list = [entry["sat_candidates"] for entry in train_ds.entries]
|
||
# Backward compat: `use_mutex_sampler=False` overrides to plain shuffle.
|
||
effective_sampler_type = cfg.sampler_type if cfg.use_mutex_sampler else "none"
|
||
|
||
if effective_sampler_type == "dss":
|
||
batch_sampler = DynamicSimilaritySampler(
|
||
sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed,
|
||
knn_device=cfg.dss_knn_device,
|
||
use_lsh=cfg.dss_use_lsh,
|
||
lsh_num_tables=cfg.dss_lsh_num_tables,
|
||
lsh_num_bits=cfg.dss_lsh_num_bits,
|
||
)
|
||
LOGGER.info(
|
||
"Sampler: DynamicSimilarity — kNN on %s%s, warmup=%d, re-embed every %d epochs",
|
||
cfg.dss_knn_device,
|
||
" + LSH" if cfg.dss_use_lsh else "",
|
||
cfg.dss_warmup_epochs, cfg.dss_reembed_every,
|
||
)
|
||
elif effective_sampler_type == "mutex":
|
||
batch_sampler = MutuallyExclusiveSampler(
|
||
sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed,
|
||
)
|
||
LOGGER.info("Sampler: MutuallyExclusive — no false negatives within a batch")
|
||
else:
|
||
batch_sampler = None
|
||
LOGGER.info("Sampler: default shuffle (no mutex / no DSS)")
|
||
|
||
if batch_sampler is not None:
|
||
train_loader = DataLoader(
|
||
train_ds,
|
||
batch_sampler=batch_sampler,
|
||
num_workers=cfg.num_workers,
|
||
collate_fn=collate_gtauav_batch,
|
||
pin_memory=True,
|
||
)
|
||
else:
|
||
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,
|
||
)
|
||
|
||
emb_cache: EmbeddingCache | None = None
|
||
if cfg.dss_cache_dir is not None:
|
||
emb_cache = EmbeddingCache(cfg.dss_cache_dir)
|
||
LOGGER.info("DSS embedding cache: %s", cfg.dss_cache_dir)
|
||
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,
|
||
)
|
||
# Train eval loader: clean transforms (no augmentation), for R@K on train set.
|
||
train_eval_ds = GTAUAVDataset(
|
||
pair_json=cfg.train_json,
|
||
rgb_root=cfg.rgb_root,
|
||
caption_root=cfg.caption_root,
|
||
image_transform=eval_tf,
|
||
filter_meta=cfg.filter_meta,
|
||
)
|
||
train_eval_loader = DataLoader(
|
||
train_eval_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,
|
||
stripnet_backbone_lr_factor=cfg.stripnet_backbone_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()
|
||
if batch_sampler is not None:
|
||
batch_sampler.set_epoch(epoch)
|
||
|
||
# DSS re-embedding: refresh query embeddings before the epoch starts.
|
||
if (
|
||
isinstance(batch_sampler, DynamicSimilaritySampler)
|
||
and epoch >= cfg.dss_warmup_epochs
|
||
and (epoch - cfg.dss_warmup_epochs) % cfg.dss_reembed_every == 0
|
||
):
|
||
query_embs: torch.Tensor | None = None
|
||
if emb_cache is not None:
|
||
query_embs = emb_cache.load(epoch)
|
||
if query_embs is None:
|
||
LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(train_ds), epoch)
|
||
t_embed = time.time()
|
||
query_embs = _embed_drone_queries(
|
||
model, train_ds, cfg.device,
|
||
batch_size=cfg.batch_size * cfg.grad_accum_steps,
|
||
num_workers=cfg.num_workers,
|
||
)
|
||
LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed)
|
||
if emb_cache is not None:
|
||
emb_cache.save(epoch, query_embs)
|
||
t_sampler = time.time()
|
||
batch_sampler.update_embeddings(query_embs)
|
||
LOGGER.info("DSS: sampler update_embeddings took %.2fs", time.time() - t_sampler)
|
||
|
||
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 — InfoNCE or WeightedInfoNCE. Only the latter uses positive_weights.
|
||
queue_neg = neg_bank.get_queue() if neg_bank is not None else None
|
||
loss_kwargs = {
|
||
"embeddings": embeddings,
|
||
"epoch": epoch,
|
||
"total_epochs": cfg.epochs,
|
||
"queue_negatives": queue_neg,
|
||
}
|
||
if isinstance(loss_fn, WeightedInfoNCELoss):
|
||
loss_kwargs["positive_weights"] = batch["positive_weights"].to(
|
||
cfg.device, non_blocking=True,
|
||
)
|
||
loss_dict = loss_fn(**loss_kwargs)
|
||
|
||
# Scale loss by accumulation steps so gradients average correctly.
|
||
raw_loss = float(loss_dict["total"].item()) # save before backward
|
||
total_loss = loss_dict["total"] / accum
|
||
scaler.scale(total_loss).backward()
|
||
|
||
# Enqueue current gallery AFTER backward. The queue buffer is aliased
|
||
# into the autograd graph through `queue_neg` (a view returned by
|
||
# `NegativeMemoryBank.get_queue`), so modifying it before backward
|
||
# triggers "variable needed for gradient computation has been modified
|
||
# by an inplace operation". Enqueueing here is semantically identical
|
||
# — the next step's queue state is the same either way.
|
||
if neg_bank is not None:
|
||
neg_bank.enqueue(embeddings["gallery"].detach())
|
||
|
||
# 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) ---
|
||
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"{step_metrics['temperature']:.4f}",
|
||
gq=f"{step_metrics['gate_q']:.3f}",
|
||
gg=f"{step_metrics['gate_g']:.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 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.
|
||
train_recall = {}
|
||
if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
|
||
# Train R@K (subset — same size as test set for speed).
|
||
train_eval_batches = len(test_loader)
|
||
train_recall = _evaluate(
|
||
model, train_eval_loader, cfg.device,
|
||
loss_fn=loss_fn, epoch=epoch, total_epochs=cfg.epochs,
|
||
max_batches=train_eval_batches, desc="eval-train",
|
||
)
|
||
epoch_record["train_recall"] = train_recall
|
||
csv_logger.log_train_recall(epoch, train_recall)
|
||
tracker.log_train(epoch, {f"recall/{k}": v for k, v in train_recall.items() if k.startswith("r@")}, step=global_step)
|
||
|
||
# Log train metrics to CSV (includes recall/AP if eval ran this epoch).
|
||
train_row = {**means}
|
||
if "total" in train_row:
|
||
train_row["train_loss"] = train_row.pop("total")
|
||
if train_recall:
|
||
train_row["r@1_q2g"] = train_recall.get("r@1_q2g", 0.0)
|
||
train_row["r@5_q2g"] = train_recall.get("r@5_q2g", 0.0)
|
||
train_row["r@10_q2g"] = train_recall.get("r@10_q2g", 0.0)
|
||
train_row["ap_q2g"] = train_recall.get("ap_q2g", 0.0)
|
||
train_row["r@1_g2q"] = train_recall.get("r@1_g2q", 0.0)
|
||
train_row["r@5_g2q"] = train_recall.get("r@5_g2q", 0.0)
|
||
train_row["r@10_g2q"] = train_recall.get("r@10_g2q", 0.0)
|
||
train_row["ap_g2q"] = train_recall.get("ap_g2q", 0.0)
|
||
csv_logger.log_train(epoch, train_row, optimizer.param_groups[0]["lr"], elapsed)
|
||
generate_plots(csv_logger.log_dir)
|
||
|
||
if train_recall:
|
||
LOGGER.info(
|
||
"train-recall epoch=%d q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f "
|
||
"g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f",
|
||
epoch,
|
||
train_recall.get("r@1_q2g", 0.0),
|
||
train_recall.get("r@5_q2g", 0.0),
|
||
train_recall.get("r@10_q2g", 0.0),
|
||
train_recall.get("ap_q2g", 0.0),
|
||
train_recall.get("r@1_g2q", 0.0),
|
||
train_recall.get("r@5_g2q", 0.0),
|
||
train_recall.get("r@10_g2q", 0.0),
|
||
train_recall.get("ap_g2q", 0.0),
|
||
train_recall.get("loss", 0.0),
|
||
)
|
||
|
||
# Val R@K (full test set).
|
||
val_metrics = _evaluate(
|
||
model, test_loader, cfg.device,
|
||
loss_fn=loss_fn, epoch=epoch, total_epochs=cfg.epochs,
|
||
desc="eval-val",
|
||
)
|
||
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 q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f "
|
||
"g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f gate_q=%.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("ap_q2g", 0.0),
|
||
val_metrics.get("r@1_g2q", 0.0),
|
||
val_metrics.get("r@5_g2q", 0.0),
|
||
val_metrics.get("r@10_g2q", 0.0),
|
||
val_metrics.get("ap_g2q", 0.0),
|
||
val_metrics.get("loss", 0.0),
|
||
val_metrics.get("gate_q", 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. Model architecture flags go into the ckpt so
|
||
# `AsymmetricEncoder.load_checkpoint` can rebuild the right shape.
|
||
_atomic_save(
|
||
obj={
|
||
"epoch": epoch,
|
||
"model_state": model.state_dict(),
|
||
"optimizer_state": optimizer.state_dict(),
|
||
"loss_state": loss_fn.state_dict(),
|
||
"baseline_mode": cfg.baseline_mode,
|
||
"shared_encoder": cfg.shared_encoder,
|
||
"mona_bottleneck": cfg.mona_bottleneck,
|
||
"mona_last_n_blocks": cfg.mona_last_n_blocks,
|
||
},
|
||
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,
|
||
loss_fn=loss_fn, epoch=cfg.epochs - 1, total_epochs=cfg.epochs,
|
||
)
|
||
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_ap_q2g": final_metrics.get("ap_q2g", 0.0),
|
||
"final_r@1_g2q": final_metrics.get("r@1_g2q", 0.0),
|
||
"final_r@5_g2q": final_metrics.get("r@5_g2q", 0.0),
|
||
"final_r@10_g2q": final_metrics.get("r@10_g2q", 0.0),
|
||
"final_ap_g2q": final_metrics.get("ap_g2q", 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 — q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f "
|
||
"g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.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("ap_q2g", 0.0),
|
||
final_metrics.get("r@1_g2q", 0.0),
|
||
final_metrics.get("r@5_g2q", 0.0),
|
||
final_metrics.get("r@10_g2q", 0.0),
|
||
final_metrics.get("ap_g2q", 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()
|