From 4e148a29bbe89983b4529f7e9cb9522a6eb28bee Mon Sep 17 00:00:00 2001 From: pikaliov Date: Mon, 4 May 2026 16:32:22 +0300 Subject: [PATCH] checkpoint 04_05 --- src/training/train.py | 62 -- src/training/train_gtauav.py | 1338 ---------------------------- src/training/train_gtauav_old.py | 1395 ------------------------------ 3 files changed, 2795 deletions(-) delete mode 100644 src/training/train.py delete mode 100644 src/training/train_gtauav.py delete mode 100644 src/training/train_gtauav_old.py diff --git a/src/training/train.py b/src/training/train.py deleted file mode 100644 index aecfac1..0000000 --- a/src/training/train.py +++ /dev/null @@ -1,62 +0,0 @@ -from __future__ import annotations - -"""Thin wrapper around src.training.trainer.Trainer. - -Kept for backward compatibility with src/main.py imports and any external -scripts that still call `train(...)` directly. After step 4b the body of -this module is just a delegation. - -Note: this module no longer runs standalone — entry point is src/main.py -(per REQUIREMENTS_GIN_STYLE.md §5): - python -m src.main -""" - -from src.conf.hardware_conf import HardwareConfig -from src.conf.models_common_conf import ModelsCommonConfig -from src.conf.models_dinov3_conf import DINOv3ModelsConfig -from src.conf.models_stripnet_conf import StripNetModelsConfig -from src.conf.pipeline_conf import PipelineConfig -from src.conf.tracking_conf import TrackingConfig -from src.conf.training_conf import TrainingConfig -from src.training.trainer_new import Trainer - -# Type alias re-exported for callers. -# SOFIA v1/v71 model configs exist in src/conf/ but are not yet supported -# by the trainer (no caption-aware fusion encoder wrapper). Adding them -# here will result in a NotImplementedError at Trainer.run(). -ModelsConfig = DINOv3ModelsConfig | StripNetModelsConfig - - -def train( - pipeline_cfg: PipelineConfig, - hardware_cfg: HardwareConfig, - training_cfg: TrainingConfig, - tracking_cfg: TrackingConfig, - models_common_cfg: ModelsCommonConfig, - models_cfg: ModelsConfig, -) -> None: - """Build a Trainer and run a full training cycle. - - Args: - pipeline_cfg: Paths, schedule (epochs/eval_every/warmup), seed, output_dir. - hardware_cfg: batch_size, grad_accum, num_workers, AMP, gradient_checkpointing. - training_cfg: Loss + optimizer + sampler recipe. - tracking_cfg: W&B / TensorBoard / Grad-CAM / profiler. - models_common_cfg: backbone, baseline_mode, init_gate, lrsclip_path. - models_cfg: Family-specific config selected by models_common_cfg.backbone. - """ - Trainer( - pipeline_cfg=pipeline_cfg, - hardware_cfg=hardware_cfg, - training_cfg=training_cfg, - tracking_cfg=tracking_cfg, - models_common_cfg=models_common_cfg, - models_cfg=models_cfg, - ).run() - - -if __name__ == "__main__": - raise SystemExit( - "Direct execution removed. Use: python -m src.main ", - ) - diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py deleted file mode 100644 index a3163b8..0000000 --- a/src/training/train_gtauav.py +++ /dev/null @@ -1,1338 +0,0 @@ -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 (via src.conf), W&B, TensorBoard, Grad-CAM, gradient -monitoring, PyTorch Profiler, and torchinfo model summary. - -Note: this module no longer runs standalone. Entry point is src/main.py -(REQUIREMENTS_GIN_STYLE.md §5): - python -m src.main -""" - -import json -import logging -import math -import time -import warnings -from pathlib import Path - -import coloredlogs -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.conf.hardware_conf import HardwareConfig -from src.conf.models_common_conf import ModelsCommonConfig -from src.conf.models_dinov3_conf import DINOv3ModelsConfig -from src.conf.models_sofia_v1_conf import SOFIAv1ModelsConfig -from src.conf.models_sofia_v71_conf import SOFIAv71ModelsConfig -from src.conf.models_stripnet_conf import StripNetModelsConfig -from src.conf.pipeline_conf import PipelineConfig -from src.conf.tracking_conf import TrackingConfig -from src.conf.training_conf import TrainingConfig -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, -) -from src.models.sofia_fusion_encoder import SOFIAFusionEncoder -from src.models.sofia_v1 import SOFIAv1Config -from src.models.sofia_v1_fusion_encoder import SOFIAv1FusionEncoder -from src.models.sofia_v71 import SOFIAConfig - -LOGGER = logging.getLogger("caption_test.train_gtauav") - -# Type alias for the family-specific models config. -ModelsConfig = ( - DINOv3ModelsConfig - | StripNetModelsConfig - | SOFIAv1ModelsConfig - | SOFIAv71ModelsConfig -) - - -def _set_seed(seed: int) -> None: - """Fix all RNG seeds for reproducibility. - - Note: duplicates src.utils.seed_utils.set_seed. Will be removed in step 4b - when this module gets decomposed into Trainer. - """ - 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: - """Save checkpoint atomically via .tmp + replace. - - Note: will be replaced with src.utils.io_utils.atomic_save_torch in step 4b - (current version doesn't clean up .tmp on error). - """ - 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: nn.Module, - lr: float, - text_lr_factor: float, - stripnet_backbone_lr_factor: float = 0.1, -) -> list[dict]: - """Build parameter groups with separate LR for text encoder and (optionally) StripNet backbone. - - Group 0: projections + heads + MONA + (logit_scale appended later). - Group 1: DGTRS-CLIP text encoder (lr * text_lr_factor). - Group 2 (optional): StripNet backbone when unfrozen (lr * stripnet_backbone_lr_factor). - """ - main_params: list[nn.Parameter] = [] - text_params: list[nn.Parameter] = [] - stripnet_backbone_params: list[nn.Parameter] = [] - - for name, p in model.named_parameters(): - if not p.requires_grad: - continue - if "text_encoder" in name: - text_params.append(p) - elif name.startswith("backbone.") or name.startswith("stripnet."): - # StripNet backbone params (only present when stripnet_freeze=False). - stripnet_backbone_params.append(p) - else: - main_params.append(p) - - groups: list[dict] = [ - {"params": main_params, "lr": lr, "name": "main"}, - ] - if text_params: - groups.append({"params": text_params, "lr": lr * text_lr_factor, "name": "text"}) - if stripnet_backbone_params: - groups.append({ - "params": stripnet_backbone_params, - "lr": lr * stripnet_backbone_lr_factor, - "name": "stripnet_backbone", - }) - return groups - - -def _cosine_warmup_schedule(warmup_steps: int, total_steps: int): - """Return a lr_lambda for LambdaLR: linear warmup + cosine decay.""" - - def lr_lambda(step: int) -> float: - if step < warmup_steps: - return float(step) / max(1, warmup_steps) - progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) - return 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0))) - - return lr_lambda - - -def _embed_drone_queries( - model: nn.Module, - train_ds, - device: str, - batch_size: int, - num_workers: int, -) -> torch.Tensor: - """Embed all drone queries from train_ds using the current model state. - - Used by DSS sampler at the start of each non-warmup epoch. - """ - model.eval() - drone_query_ds = GTAUAVDroneQuery( - train_ds.entries, - rgb_root=str(train_ds.rgb_root), - caption_root=str(train_ds.caption_root), - image_transform=get_dino_transform(image_size=256), - ) - loader = DataLoader( - drone_query_ds, - batch_size=batch_size, - shuffle=False, - num_workers=num_workers, - collate_fn=collate_drone_query, - pin_memory=True, - ) - all_embs: list[torch.Tensor] = [] - with torch.inference_mode(): - for batch in tqdm(loader, desc="dss-embed", unit="batch", leave=False): - drone_img = batch["drone_img"].to(device, non_blocking=True) - altitude = batch.get("altitude") - if altitude is not None: - altitude = altitude.to(device, non_blocking=True) - kwargs = {"drone_img": drone_img, "altitude": altitude} - if not getattr(model, "baseline_mode", False): - kwargs["caption_l1"] = batch["caption_l1"] - kwargs["caption_l2"] = batch["caption_l2"] - kwargs["caption_l3"] = batch["caption_l3"] - with autocast(device_type="cuda", enabled=True): - emb = model.encode_drone_query(**kwargs) - all_embs.append(emb.cpu()) - model.train() - return torch.cat(all_embs, dim=0) - - -@torch.no_grad() -def _evaluate( - model: nn.Module, - loader: DataLoader, - device: str, - loss_fn: nn.Module, - epoch: int, - total_epochs: int, - 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) - altitude = batch.get("altitude") - if altitude is not None: - altitude = altitude.to(device, non_blocking=True) - q = model.encode_query( - drone_img, - batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], - altitude=altitude, - ) - 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 and reset peak memory stats. - - Note: duplicates src.utils.io_utils.clear_vram. Will be replaced in step 4b. - """ - import gc as _gc - _gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.reset_peak_memory_stats() - allocated_gb = torch.cuda.memory_allocated() / 1e9 - LOGGER.info("VRAM cleared. Current usage: %.2f GB", allocated_gb) - - -def train( - pipeline_cfg: PipelineConfig, - hardware_cfg: HardwareConfig, - training_cfg: TrainingConfig, - tracking_cfg: TrackingConfig, - models_common_cfg: ModelsCommonConfig, - models_cfg: ModelsConfig, -) -> None: - """Run full training loop. - - Args: - pipeline_cfg: Paths, schedule (epochs/eval_every/warmup), seed, output_dir. - hardware_cfg: batch_size, grad_accum, num_workers, AMP, gradient_checkpointing. - training_cfg: Loss + optimizer + sampler recipe. - tracking_cfg: W&B / TensorBoard / Grad-CAM / profiler. - models_common_cfg: backbone, baseline_mode, init_gate, lrsclip_path. - models_cfg: Family-specific config selected by models_common_cfg.backbone: - DINOv3ModelsConfig | StripNetModelsConfig - | SOFIAv1ModelsConfig | SOFIAv71ModelsConfig. - """ - coloredlogs.install( - level="INFO", - logger=LOGGER, - fmt="%(asctime)s %(name)s %(levelname)s %(message)s", - ) - _clear_vram() - _set_seed(pipeline_cfg.seed) - output_dir = Path(pipeline_cfg.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - # Save config — all 6 config objects merged into one dict for traceability. - full_config = { - "pipeline": vars(pipeline_cfg), - "hardware": vars(hardware_cfg), - "training": vars(training_cfg), - "tracking": vars(tracking_cfg), - "models_common": vars(models_common_cfg), - "models": vars(models_cfg), - } - with (output_dir / "config.json").open("w") as f: - json.dump(full_config, f, indent=2) - - # --- Experiment tracker (W&B + TensorBoard) --- - tracker = ExperimentTracker( - output_dir=output_dir, - config=full_config, - use_wandb=tracking_cfg.use_wandb, - use_tb=tracking_cfg.use_tb, - wandb_project=tracking_cfg.wandb_project, - wandb_run_name=tracking_cfg.wandb_run_name, - wandb_entity=tracking_cfg.wandb_entity, - ) - - # Model. - backbone = models_common_cfg.backbone - start_epoch = 0 - resume_ckpt = None - - if pipeline_cfg.resume_from is not None: - LOGGER.info("Resuming from %s", pipeline_cfg.resume_from) - if backbone == "sofia_v71": - model, resume_ckpt = SOFIAFusionEncoder.load_checkpoint( - pipeline_cfg.resume_from, - lrsclip_path=models_common_cfg.lrsclip_path, - device=hardware_cfg.device, - ) - elif backbone == "sofia_v1": - model, resume_ckpt = SOFIAv1FusionEncoder.load_checkpoint( - pipeline_cfg.resume_from, - lrsclip_path=models_common_cfg.lrsclip_path, - device=hardware_cfg.device, - ) - else: - # DINOv3 or StripNet — both go through AsymmetricEncoder. - assert isinstance(models_cfg, (DINOv3ModelsConfig, StripNetModelsConfig)), ( - f"Expected DINOv3/StripNet ModelsConfig for backbone={backbone!r}, " - f"got {type(models_cfg).__name__}" - ) - dino_web_path = ( - models_cfg.dino_web_path if isinstance(models_cfg, DINOv3ModelsConfig) - else "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth" - ) - dino_sat_path = ( - models_cfg.dino_sat_path if isinstance(models_cfg, DINOv3ModelsConfig) - else "nn_models/DINO_SAT/model.safetensors" - ) - model, resume_ckpt = AsymmetricEncoder.load_checkpoint( - pipeline_cfg.resume_from, - dino_web_path=dino_web_path, - dino_sat_path=dino_sat_path, - lrsclip_path=models_common_cfg.lrsclip_path, - device=hardware_cfg.device, - ) - start_epoch = resume_ckpt.get("epoch", -1) + 1 - else: - mode_str = "baseline (no text)" if models_common_cfg.baseline_mode else "with text (L1/L2/L3)" - if backbone == "sofia_v71": - assert isinstance(models_cfg, SOFIAv71ModelsConfig) - enc_str = ( - f"SOFIA-{models_cfg.variant_label} " - f"(text-FiLM uav={models_cfg.use_text_film_uav}, " - f"sat={models_cfg.use_text_film_sat})" - ) - elif backbone == "sofia_v1": - assert isinstance(models_cfg, SOFIAv1ModelsConfig) - enc_str = ( - f"SOFIAv1-{models_cfg.variant_label} (StripNet+DCNv4, " - f"text-FiLM uav={models_cfg.use_text_film_uav}, " - f"sat={models_cfg.use_text_film_sat})" - ) - elif backbone == "stripnet": - enc_str = "StripNet-small (shared, 512→1024 proj)" - else: # dinov3 - assert isinstance(models_cfg, DINOv3ModelsConfig) - enc_str = "shared DINOv3 WEB" if models_cfg.shared_encoder else "asymmetric (WEB + SAT)" - LOGGER.info("Building model — %s, %s", mode_str, enc_str) - - if backbone == "sofia_v71": - assert isinstance(models_cfg, SOFIAv71ModelsConfig) - # Build SOFIAConfig from the gin-loaded SOFIAv71ModelsConfig. - # All architectural fields come from gin — no preset factory needed. - sofia_cfg = SOFIAConfig( - input_size=models_cfg.input_size, - in_channels=models_cfg.in_channels, - stem_mid=models_cfg.stem_mid, - stem_out=models_cfg.stem_out, - embed_dims=list(models_cfg.embed_dims), - depths=list(models_cfg.depths), - mbconv_expand=models_cfg.mbconv_expand, - se_ratio=models_cfg.se_ratio, - strip_kernel_s1=models_cfg.strip_kernel_s1, - strip_kernel_s2=models_cfg.strip_kernel_s2, - mix_kernels=list(models_cfg.mix_kernels), - use_dcn_strip=models_cfg.use_dcn_strip, - mamba_d_state=models_cfg.mamba_d_state, - mamba_dt_rank=models_cfg.mamba_dt_rank, - mamba_backend=models_cfg.mamba_backend, - mamba_variant=models_cfg.mamba_variant, - mamba_extra_kwargs=dict(models_cfg.mamba_extra_kwargs), - num_heads_s3=models_cfg.num_heads_s3, - num_heads_s4=models_cfg.num_heads_s4, - use_strip_branch_s3=models_cfg.use_strip_branch_s3, - use_strip_branch_s4=models_cfg.use_strip_branch_s4, - ffn_expand=models_cfg.ffn_expand, - use_evss_bridge=models_cfg.use_evss_bridge, - evss_bridge_locations=list(models_cfg.evss_bridge_locations), - neck_channels=models_cfg.neck_channels, - d_descriptor=models_cfg.d_descriptor, - use_asymmetric_heads=models_cfg.use_asymmetric_heads, - chp_rings=models_cfg.chp_rings, - chp_angles=models_cfg.chp_angles, - chp_harmonics=models_cfg.chp_harmonics, - use_film_altitude=models_cfg.use_film_altitude, - altitude_norm=models_cfg.altitude_norm, - ring_count=models_cfg.ring_count, - use_ring_aux=models_cfg.use_ring_aux, - return_normalized=models_cfg.return_normalized, - # Disable text fusion when baseline_mode is on, regardless of gin. - use_text_film_sat=models_cfg.use_text_film_sat and not models_common_cfg.baseline_mode, - use_text_film_uav=models_cfg.use_text_film_uav and not models_common_cfg.baseline_mode, - text_film_dim=models_cfg.text_film_dim, - text_film_hidden=models_cfg.text_film_hidden, - share_stages_1_2=models_cfg.share_stages_1_2, - enable_kd_taps=models_cfg.enable_kd_taps, - precision=models_cfg.precision, - ) - model = SOFIAFusionEncoder( - sofia_cfg=sofia_cfg, - lrsclip_path=models_common_cfg.lrsclip_path, - init_gate=models_common_cfg.init_gate, - baseline_mode=models_common_cfg.baseline_mode, - lora_rank=models_cfg.lora_rank, - device=hardware_cfg.device, - ).to(hardware_cfg.device) - elif backbone == "sofia_v1": - assert isinstance(models_cfg, SOFIAv1ModelsConfig) - sofia_v1_cfg = SOFIAv1Config( - variant=models_cfg.variant_label, - dcn_variant=models_cfg.dcn_variant, - d_descriptor=models_cfg.d_descriptor, - text_film_dim=models_cfg.d_descriptor, # match d_descriptor (preserves old behavior) - use_text_film_uav=models_cfg.use_text_film_uav and not models_common_cfg.baseline_mode, - use_text_film_sat=models_cfg.use_text_film_sat and not models_common_cfg.baseline_mode, - use_film_altitude=models_cfg.use_film_altitude, - ) - model = SOFIAv1FusionEncoder( - sofia_cfg=sofia_v1_cfg, - lrsclip_path=models_common_cfg.lrsclip_path, - init_gate=models_common_cfg.init_gate, - baseline_mode=models_common_cfg.baseline_mode, - lora_rank=models_cfg.lora_rank, - device=hardware_cfg.device, - ).to(hardware_cfg.device) - elif backbone == "stripnet": - assert isinstance(models_cfg, StripNetModelsConfig) - # AsymmetricEncoder also handles StripNet — pass dummy DINO paths, - # they're not used when backbone='stripnet'. (DINO fields not - # bindable on StripNetModelsConfig — by design.) - model = AsymmetricEncoder( - dino_web_path="nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth", - dino_sat_path="nn_models/DINO_SAT/model.safetensors", - lrsclip_path=models_common_cfg.lrsclip_path, - init_gate=models_common_cfg.init_gate, - baseline_mode=models_common_cfg.baseline_mode, - shared_encoder=True, # StripNet is always shared - mona_bottleneck=64, - mona_last_n_blocks=12, - device=hardware_cfg.device, - backbone=backbone, - stripnet_path=models_cfg.stripnet_path, - stripnet_mona_last_n_stages=models_cfg.stripnet_mona_last_n_stages, - stripnet_freeze=models_cfg.stripnet_freeze, - ).to(hardware_cfg.device) - else: # dinov3 - assert isinstance(models_cfg, DINOv3ModelsConfig) - model = AsymmetricEncoder( - dino_web_path=models_cfg.dino_web_path, - dino_sat_path=models_cfg.dino_sat_path, - lrsclip_path=models_common_cfg.lrsclip_path, - init_gate=models_common_cfg.init_gate, - baseline_mode=models_common_cfg.baseline_mode, - shared_encoder=models_cfg.shared_encoder, - mona_bottleneck=models_cfg.mona_bottleneck, - mona_last_n_blocks=models_cfg.mona_last_n_blocks, - device=hardware_cfg.device, - backbone=backbone, - stripnet_path="nn_models/STRIPNET/stripnet_s.pth", - stripnet_mona_last_n_stages=0, - stripnet_freeze=True, - ).to(hardware_cfg.device) - LOGGER.info("embed_dim=%d", model.embed_dim) - - # --- Gradient checkpointing (trade compute for VRAM) --- - # StripNet/SOFIA don't expose set_gradient_checkpointing — only DGTRS gets it. - if hardware_cfg.gradient_checkpointing and backbone == "dinov3": - assert isinstance(models_cfg, DINOv3ModelsConfig) - if models_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 hardware_cfg.gradient_checkpointing and backbone in ("stripnet", "sofia_v71", "sofia_v1"): - if model.text_encoder is not None: - model.text_encoder.transformer.gradient_checkpointing = True - LOGGER.info("Gradient checkpointing enabled (DGTRS only; %s doesn't support)", backbone) - - 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=hardware_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. InfoNCELoss / WeightedInfoNCELoss are NOT @gin.configurable — - # all parameters arrive explicitly from training_cfg. - if training_cfg.loss_type == "symmetric": - loss_fn = InfoNCELoss( - temperature_init=training_cfg.tau_init, - temperature_final=training_cfg.tau_final, - label_smoothing=training_cfg.label_smoothing, - weight_q2g=training_cfg.weight_q2g, - weight_g2q=training_cfg.weight_g2q, - learnable_temperature=training_cfg.learnable_temperature, - tau_min=training_cfg.tau_min, - tau_max=training_cfg.tau_max, - hard_mining_k=training_cfg.hard_mining_k, - ) - loss_name = "SymmetricInfoNCE" - elif training_cfg.loss_type == "weighted": - loss_fn = WeightedInfoNCELoss( - temperature_init=training_cfg.tau_init, - learnable_temperature=training_cfg.learnable_temperature, - label_smoothing=training_cfg.label_smoothing, - k=training_cfg.weighted_loss_k, - tau_min=training_cfg.tau_min, - tau_max=training_cfg.tau_max, - ) - loss_name = "WeightedInfoNCE" - else: - raise ValueError( - f"Unknown loss_type={training_cfg.loss_type!r} " - f"(expected 'symmetric' or 'weighted')", - ) - - LOGGER.info( - "Loss: %s Temperature: %s (init=%.3f) q2g=%.2f g2q=%.2f", - loss_name, - "learnable" if training_cfg.learnable_temperature else "fixed", - training_cfg.tau_init, training_cfg.weight_q2g, training_cfg.weight_g2q, - ) - - # Hard negative memory bank. - neg_bank = None - if training_cfg.neg_bank_size > 0: - neg_bank = NegativeMemoryBank(size=training_cfg.neg_bank_size, dim=model.embed_dim).to(hardware_cfg.device) - LOGGER.info("Negative memory bank: size=%d, dim=%d", training_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=pipeline_cfg.train_json, - rgb_root=pipeline_cfg.rgb_root, - caption_root=pipeline_cfg.caption_root, - drone_transform=drone_train_tf, - sat_transform=sat_train_tf, - filter_meta=pipeline_cfg.filter_meta, - ) - test_ds = GTAUAVDataset( - pair_json=pipeline_cfg.test_json, - rgb_root=pipeline_cfg.rgb_root, - caption_root=pipeline_cfg.caption_root, - image_transform=eval_tf, - filter_meta=pipeline_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 = training_cfg.sampler_type if training_cfg.use_mutex_sampler else "none" - - if effective_sampler_type == "dss": - batch_sampler = DynamicSimilaritySampler( - sat_cand_list, - batch_size=hardware_cfg.batch_size, - shuffle=True, - seed=pipeline_cfg.seed, - knn_device=training_cfg.dss_knn_device, - use_lsh=training_cfg.dss_use_lsh, - lsh_num_tables=training_cfg.dss_lsh_num_tables, - lsh_num_bits=training_cfg.dss_lsh_num_bits, - ) - LOGGER.info( - "Sampler: DynamicSimilarity — kNN on %s%s, warmup=%d, re-embed every %d epochs", - training_cfg.dss_knn_device, - " + LSH" if training_cfg.dss_use_lsh else "", - training_cfg.dss_warmup_epochs, training_cfg.dss_reembed_every, - ) - elif effective_sampler_type == "mutex": - batch_sampler = MutuallyExclusiveSampler( - sat_cand_list, - batch_size=hardware_cfg.batch_size, - shuffle=True, - seed=pipeline_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=hardware_cfg.num_workers, - collate_fn=collate_gtauav_batch, - pin_memory=True, - ) - else: - train_loader = DataLoader( - train_ds, - batch_size=hardware_cfg.batch_size, - shuffle=True, - num_workers=hardware_cfg.num_workers, - collate_fn=collate_gtauav_batch, - pin_memory=True, - drop_last=True, - ) - - emb_cache: EmbeddingCache | None = None - if training_cfg.dss_cache_dir is not None: - emb_cache = EmbeddingCache(training_cfg.dss_cache_dir) - LOGGER.info("DSS embedding cache: %s", training_cfg.dss_cache_dir) - test_loader = DataLoader( - test_ds, - batch_size=hardware_cfg.batch_size, - shuffle=False, - num_workers=hardware_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=pipeline_cfg.train_json, - rgb_root=pipeline_cfg.rgb_root, - caption_root=pipeline_cfg.caption_root, - image_transform=eval_tf, - filter_meta=pipeline_cfg.filter_meta, - ) - train_eval_loader = DataLoader( - train_eval_ds, - batch_size=hardware_cfg.batch_size, - shuffle=False, - num_workers=hardware_cfg.num_workers, - collate_fn=collate_gtauav_batch, - pin_memory=True, - ) - - effective_batch = hardware_cfg.batch_size * hardware_cfg.grad_accum_steps - LOGGER.info( - "train=%d test=%d batch=%d accum=%d effective_batch=%d", - len(train_ds), len(test_ds), - hardware_cfg.batch_size, hardware_cfg.grad_accum_steps, effective_batch, - ) - - # Optimizer — per-group LR (text encoder gets lower LR, StripNet backbone optionally). - stripnet_lr_factor = ( - models_cfg.stripnet_backbone_lr_factor - if isinstance(models_cfg, StripNetModelsConfig) - else 0.1 # default; not used unless StripNet group is non-empty - ) - param_groups = _build_param_groups( - model, - training_cfg.learning_rate, - training_cfg.text_lr_factor, - stripnet_backbone_lr_factor=stripnet_lr_factor, - ) - # Include loss temperature if learnable. - if training_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=training_cfg.weight_decay) - - lr_info = f"proj={training_cfg.learning_rate:.0e}" - if not models_common_cfg.baseline_mode: - lr_info += f" text={training_cfg.learning_rate * training_cfg.text_lr_factor:.0e}" - LOGGER.info("Optimizer: AdamW LR: %s warmup=%d epochs", lr_info, pipeline_cfg.warmup_epochs) - - # Scheduler — cosine with linear warmup (counted in optimizer steps). - steps_per_epoch = math.ceil(len(train_loader) / hardware_cfg.grad_accum_steps) - total_steps = pipeline_cfg.epochs * steps_per_epoch - warmup_steps = pipeline_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=hardware_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 tracking_cfg.use_profiler and start_epoch == 0: - profiler = TrainingProfiler( - output_dir=output_dir, - n_warmup=tracking_cfg.profiler_warmup, - n_active=tracking_cfg.profiler_active, - ) - profiler.start() - - LOGGER.info("Starting training for %d epochs (from epoch %d)", pipeline_cfg.epochs, start_epoch) - - global_step = start_epoch * steps_per_epoch - best_r1 = 0.0 - - for epoch in range(start_epoch, pipeline_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 >= training_cfg.dss_warmup_epochs - and (epoch - training_cfg.dss_warmup_epochs) % training_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, hardware_cfg.device, - batch_size=hardware_cfg.batch_size * hardware_cfg.grad_accum_steps, - num_workers=hardware_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}/{pipeline_cfg.epochs}", - unit="batch", - leave=False, - ) - accum = hardware_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(hardware_cfg.device, non_blocking=True) - sat_img = batch["sat_img"].to(hardware_cfg.device, non_blocking=True) - altitude = batch.get("altitude") - if altitude is not None: - altitude = altitude.to(hardware_cfg.device, non_blocking=True) - - # Model forward in AMP (fp16 for DINOv3/DGTRS encoders). - with autocast(device_type="cuda", enabled=hardware_cfg.use_amp): - if models_common_cfg.baseline_mode: - embeddings = model(drone_img=drone_img, sat_img=sat_img, altitude=altitude) - 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"], - altitude=altitude, - ) - # 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": pipeline_cfg.epochs, - "queue_negatives": queue_neg, - } - if isinstance(loss_fn, WeightedInfoNCELoss): - loss_kwargs["positive_weights"] = batch["positive_weights"].to( - hardware_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. - 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 training_cfg.grad_clip > 0: - scaler.unscale_(optimizer) - nn.utils.clip_grad_norm_( - model.trainable_parameters(), - max_norm=training_cfg.grad_clip, - ) - - # --- Gradient monitoring (after unscale, before step) --- - if tracking_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) % pipeline_cfg.eval_every == 0 or epoch == pipeline_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, hardware_cfg.device, - loss_fn=loss_fn, epoch=epoch, total_epochs=pipeline_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, hardware_cfg.device, - loss_fn=loss_fn, epoch=epoch, total_epochs=pipeline_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 tracking_cfg.use_gradcam and (epoch + 1) % tracking_cfg.gradcam_every == 0: - from src.training.gradcam import generate_gradcam_samples - overlays = generate_gradcam_samples( - model=model, - dataloader=test_loader, - device=hardware_cfg.device, - output_dir=str(output_dir), - n_samples=tracking_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` (or `SOFIAFusionEncoder.load_checkpoint`) - # can rebuild the right shape. - ckpt_obj = { - "epoch": epoch, - "model_state": model.state_dict(), - "optimizer_state": optimizer.state_dict(), - "loss_state": loss_fn.state_dict(), - "baseline_mode": models_common_cfg.baseline_mode, - "backbone": backbone, - } - if backbone in ("sofia_v71", "sofia_v1"): - ckpt_obj["sofia_cfg"] = model.sofia_cfg - elif isinstance(models_cfg, DINOv3ModelsConfig): - ckpt_obj["shared_encoder"] = models_cfg.shared_encoder - ckpt_obj["mona_bottleneck"] = models_cfg.mona_bottleneck - ckpt_obj["mona_last_n_blocks"] = models_cfg.mona_last_n_blocks - # StripNet has no extra arch flags worth saving here (params come from gin on resume). - _atomic_save(obj=ckpt_obj, 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, hardware_cfg.device, - loss_fn=loss_fn, epoch=pipeline_cfg.epochs - 1, total_epochs=pipeline_cfg.epochs, - ) - report = { - "config": full_config, - "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), - ) - - -# Direct execution removed — entry point is src/main.py per REQUIREMENTS_GIN_STYLE.md §5. -if __name__ == "__main__": - raise SystemExit( - "Direct execution removed. Use: python -m src.main ", - ) - diff --git a/src/training/train_gtauav_old.py b/src/training/train_gtauav_old.py deleted file mode 100644 index 026a610..0000000 --- a/src/training/train_gtauav_old.py +++ /dev/null @@ -1,1395 +0,0 @@ -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, -) -from src.models.sofia_fusion_encoder import SOFIAFusionEncoder -from src.models.sofia_v1 import SOFIAv1Config -from src.models.sofia_v1_fusion_encoder import SOFIAv1FusionEncoder -from src.models.sofia_v71 import ( - sofia_l_config, - sofia_m_config, - sofia_tiny_config, -) - -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", "stripnet", or "sofia" - 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) - # SOFIA backbone options (used when backbone="sofia"). Trained from scratch — no pretrained checkpoint. - sofia_preset: str = "Tiny" # "Tiny" | "M" | "L" - sofia_d_descriptor: int = 1024 # retrieval space (1024 = match TextFusionMLP out_dim) - sofia_use_text_film_uav: bool = True # mid-level text-FiLM in UAV head - sofia_use_text_film_sat: bool = True # mid-level text-FiLM in SAT head - sofia_lora_rank: int = 4 - sofia_mamba_variant: str = "mamba2" # "mamba1" | "mamba2" | "efficient_vmamba" - sofia_mamba_backend: str = "auto" # "auto" | "torch" | "mamba_ssm" - # EVSSBridge (B6-inspired refinement between heterogeneous stages, opt-in). - sofia_use_evss_bridge: bool = False - sofia_evss_bridge_locations: list[str] = field(default_factory=lambda: ["pre_stage3"]) - # SOFIA v1 backbone options (used when backbone="sofia_v1"). StripNet+DCN, from scratch. - sofia_v1_variant: str = "small" # "tiny_tiny" | "tiny" | "small" | "small_v2" - sofia_v1_dcn_variant: str = "v2" # "v2" (torchvision DeformConv2d, stable) | "v4" (OpenGVLab, leaky) - sofia_v1_d_descriptor: int = 1024 - sofia_v1_use_text_film_uav: bool = True - sofia_v1_use_text_film_sat: bool = True - sofia_v1_use_film_altitude: bool = True - sofia_v1_lora_rank: int = 4 - - # 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. - 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: nn.Module, - 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, - SOFIA backbone+heads when backbone="sofia") → 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) - altitude = batch.get("altitude") - if altitude is not None: - altitude = altitude.to(device, non_blocking=True) - q = model.encode_query( - drone_img, - batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], - altitude=altitude, - ) - 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) - altitude = batch.get("altitude") - if altitude is not None: - altitude = altitude.to(device, non_blocking=True) - q = model.encode_query( - drone_img, - batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], - altitude=altitude, - ) - 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) - if cfg.backbone == "sofia_v71": - model, resume_ckpt = SOFIAFusionEncoder.load_checkpoint( - cfg.resume_from, - lrsclip_path=cfg.lrsclip_path, - device=cfg.device, - ) - elif cfg.backbone == "sofia_v1": - model, resume_ckpt = SOFIAv1FusionEncoder.load_checkpoint( - cfg.resume_from, - lrsclip_path=cfg.lrsclip_path, - device=cfg.device, - ) - else: - 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 == "sofia_v71": - enc_str = f"SOFIA-{cfg.sofia_preset} (text-FiLM uav={cfg.sofia_use_text_film_uav}, sat={cfg.sofia_use_text_film_sat})" - elif cfg.backbone == "sofia_v1": - enc_str = f"SOFIAv1-{cfg.sofia_v1_variant} (StripNet+DCNv4, text-FiLM uav={cfg.sofia_v1_use_text_film_uav}, sat={cfg.sofia_v1_use_text_film_sat})" - elif 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) - if cfg.backbone == "sofia_v71": - preset_map = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config} - if cfg.sofia_preset not in preset_map: - raise ValueError(f"Unknown sofia_preset={cfg.sofia_preset!r}") - sofia_cfg = preset_map[cfg.sofia_preset]() - sofia_cfg.d_descriptor = cfg.sofia_d_descriptor - sofia_cfg.text_film_dim = cfg.sofia_d_descriptor - sofia_cfg.use_text_film_uav = cfg.sofia_use_text_film_uav and not cfg.baseline_mode - sofia_cfg.use_text_film_sat = cfg.sofia_use_text_film_sat and not cfg.baseline_mode - sofia_cfg.mamba_variant = cfg.sofia_mamba_variant - sofia_cfg.mamba_backend = cfg.sofia_mamba_backend - sofia_cfg.use_evss_bridge = cfg.sofia_use_evss_bridge - sofia_cfg.evss_bridge_locations = list(cfg.sofia_evss_bridge_locations) - model = SOFIAFusionEncoder( - sofia_cfg=sofia_cfg, - lrsclip_path=cfg.lrsclip_path, - init_gate=cfg.init_gate, - baseline_mode=cfg.baseline_mode, - lora_rank=cfg.sofia_lora_rank, - device=cfg.device, - ).to(cfg.device) - elif cfg.backbone == "sofia_v1": - sofia_v1_cfg = SOFIAv1Config( - variant=cfg.sofia_v1_variant, - dcn_variant=cfg.sofia_v1_dcn_variant, - d_descriptor=cfg.sofia_v1_d_descriptor, - text_film_dim=cfg.sofia_v1_d_descriptor, - use_text_film_uav=cfg.sofia_v1_use_text_film_uav and not cfg.baseline_mode, - use_text_film_sat=cfg.sofia_v1_use_text_film_sat and not cfg.baseline_mode, - use_film_altitude=cfg.sofia_v1_use_film_altitude, - ) - model = SOFIAv1FusionEncoder( - sofia_cfg=sofia_v1_cfg, - lrsclip_path=cfg.lrsclip_path, - init_gate=cfg.init_gate, - baseline_mode=cfg.baseline_mode, - lora_rank=cfg.sofia_v1_lora_rank, - device=cfg.device, - ).to(cfg.device) - else: - 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/SOFIA don't expose set_gradient_checkpointing — only DGTRS gets it. - 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 in ("stripnet", "sofia_v71", "sofia_v1"): - if model.text_encoder is not None: - model.text_encoder.transformer.gradient_checkpointing = True - LOGGER.info("Gradient checkpointing enabled (DGTRS only; %s doesn't support)", cfg.backbone) - - 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) - altitude = batch.get("altitude") - if altitude is not None: - altitude = altitude.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, altitude=altitude) - 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"], - altitude=altitude, - ) - # 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` (or `SOFIAFusionEncoder.load_checkpoint`) - # can rebuild the right shape. - ckpt_obj = { - "epoch": epoch, - "model_state": model.state_dict(), - "optimizer_state": optimizer.state_dict(), - "loss_state": loss_fn.state_dict(), - "baseline_mode": cfg.baseline_mode, - "backbone": cfg.backbone, - } - if cfg.backbone in ("sofia_v71", "sofia_v1"): - ckpt_obj["sofia_cfg"] = model.sofia_cfg - else: - ckpt_obj["shared_encoder"] = cfg.shared_encoder - ckpt_obj["mona_bottleneck"] = cfg.mona_bottleneck - ckpt_obj["mona_last_n_blocks"] = cfg.mona_last_n_blocks - _atomic_save(obj=ckpt_obj, 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()