diff --git a/src/training/trainer_new.py b/src/training/trainer_new.py new file mode 100644 index 0000000..3fb8b59 --- /dev/null +++ b/src/training/trainer_new.py @@ -0,0 +1,1055 @@ +from __future__ import annotations + +"""Trainer for CVGL caption test on GTA-UAV-LR. + +Decomposed from src/training/train_gtauav.py::train into a class with one +orchestrating method `run()` plus dedicated `_setup_*` / `_build_*` / +`_train_*` / `_evaluate_*` methods. + +Lifecycle: + Trainer(...) → run() → done. + +`run()` calls _build_* in dependency order, then _train_loop, then +_final_evaluation; cleanup is in a `finally` block. + +Currently supports DINOv3 and StripNet backbones only. SOFIA v1/v7.1 model +classes live in src/models/sofia_v1/ and src/models/sofia_v71/ but are not +yet wired into this training pipeline (no caption-aware fusion encoder +wrapper exists for them). Their gin presets remain in in/config_files/ +for future integration; loading one will fail at config_loader level. +""" + +import json +import logging +import math +import time +import warnings +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from torch.amp import GradScaler, autocast +from torch.optim import AdamW +from torch.optim.lr_scheduler import LambdaLR +from torch.utils.data import DataLoader +from tqdm import tqdm + +from src.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.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler +from src.datasets.embedding_cache import EmbeddingCache +from src.datasets.gtauav_dataset import ( + GTAUAVDataset, + GTAUAVDroneQuery, + collate_drone_query, + collate_gtauav_batch, +) +from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler +from src.eval.evaluator import evaluate +from src.losses.hard_negatives import NegativeMemoryBank +from src.losses.multi_infonce import InfoNCELoss +from src.losses.weighted_infonce import WeightedInfoNCELoss +from src.models.asymmetric_encoder import ( + AsymmetricEncoder, + get_dino_transform, + get_drone_train_transform, + get_satellite_train_transform, +) +from src.training.csv_logger import CSVLogger +from src.training.grad_monitor import compute_gradient_norms, log_gradient_summary +from src.training.plot_metrics import generate_plots +from src.training.profiling import TrainingProfiler, print_model_summary +from src.training.trackers import ExperimentTracker +from src.utils.io_utils import atomic_save_torch, clear_vram +from src.utils.seed_utils import set_seed + +LOGGER = logging.getLogger("caption_test.trainer") + +# Type alias for the family-specific models config. +# SOFIA v1/v71 will join this union once their fusion encoders are written. +ModelsConfig = DINOv3ModelsConfig | StripNetModelsConfig + +# Backbones currently wired into this trainer. +_SUPPORTED_BACKBONES: frozenset[str] = frozenset({"dinov3", "stripnet"}) + + +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 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.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: 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, + ) + 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: dict[str, Any] = {"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_query(**kwargs) + all_embs.append(emb.cpu()) + if was_training: + model.train() + return torch.cat(all_embs, dim=0) + + +class Trainer: + """Orchestrates one training run. + + All gin parameters arrive as 6 config objects; runtime state (model, + optimizer, loaders, ...) is built lazily by _build_* methods and lives + on `self`. `run()` calls them in dependency order. + + Backbones supported: 'dinov3', 'stripnet'. + """ + + def __init__( + self, + pipeline_cfg: PipelineConfig, + hardware_cfg: HardwareConfig, + training_cfg: TrainingConfig, + tracking_cfg: TrackingConfig, + models_common_cfg: ModelsCommonConfig, + models_cfg: ModelsConfig, + ) -> None: + self.pipeline_cfg = pipeline_cfg + self.hardware_cfg = hardware_cfg + self.training_cfg = training_cfg + self.tracking_cfg = tracking_cfg + self.models_common_cfg = models_common_cfg + self.models_cfg = models_cfg + + # Runtime state — populated by _build_* methods. + self.output_dir: Path | None = None + self.full_config: dict | None = None + self.tracker: ExperimentTracker | None = None + self.csv_logger: CSVLogger | None = None + self.model: nn.Module | None = None + self.loss_fn: nn.Module | None = None + self.neg_bank: NegativeMemoryBank | None = None + self.optimizer: torch.optim.Optimizer | None = None + self.scheduler: LambdaLR | None = None + self.scaler: GradScaler | None = None + self.train_ds: GTAUAVDataset | None = None + self.test_ds: GTAUAVDataset | None = None + self.train_eval_ds: GTAUAVDataset | None = None + self.train_loader: DataLoader | None = None + self.test_loader: DataLoader | None = None + self.train_eval_loader: DataLoader | None = None + self.batch_sampler: DynamicSimilaritySampler | MutuallyExclusiveSampler | None = None + self.emb_cache: EmbeddingCache | None = None + self.profiler: TrainingProfiler | None = None + self.resume_ckpt: dict | None = None + + # Loop state. + self.start_epoch: int = 0 + self.global_step: int = 0 + self.best_r1: float = 0.0 + self.history: list[dict] = [] + self.steps_per_epoch: int = 0 + + # =================================================================== + # Public entry point + # =================================================================== + + def run(self) -> None: + """Full pipeline: setup → build → train → evaluate → cleanup.""" + self._validate_backbone() + clear_vram() + set_seed(self.pipeline_cfg.seed) + self._setup_output_dir() + self._setup_tracker() + self._build_model() + self._configure_gradient_checkpointing() + self._log_model_summary() + self._build_loss() + self._build_neg_bank() + self._build_data_loaders() + self._build_optimizer_and_scheduler() + self._restore_from_resume() + self._setup_profiler() + + try: + self._train_loop() + self._final_evaluation() + finally: + self._cleanup() + + # =================================================================== + # Build phase + # =================================================================== + + def _validate_backbone(self) -> None: + """Reject unsupported backbones up front with a helpful message.""" + backbone = self.models_common_cfg.backbone + if backbone not in _SUPPORTED_BACKBONES: + raise NotImplementedError( + f"Trainer does not support backbone={backbone!r} yet. " + f"Supported backbones: {sorted(_SUPPORTED_BACKBONES)}. " + f"SOFIA v1/v7.1 model classes exist in src/models/sofia_v1/ and " + f"src/models/sofia_v71/, but a caption-aware fusion encoder " + f"wrapper has not been written for them. To enable, create " + f"the wrapper class with .encode_query/.encode_gallery/" + f".fusion_query/.fusion_gallery API matching AsymmetricEncoder, " + f"then add the corresponding branch to _build_model.", + ) + + def _setup_output_dir(self) -> None: + """Create output_dir, save config.json, init csv_logger.""" + self.output_dir = Path(self.pipeline_cfg.output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + # Merge all 6 config objects into one dict for full traceability. + self.full_config = { + "pipeline": vars(self.pipeline_cfg), + "hardware": vars(self.hardware_cfg), + "training": vars(self.training_cfg), + "tracking": vars(self.tracking_cfg), + "models_common": vars(self.models_common_cfg), + "models": vars(self.models_cfg), + } + with (self.output_dir / "config.json").open("w") as f: + json.dump(self.full_config, f, indent=2) + + self.csv_logger = CSVLogger(self.output_dir) + + def _setup_tracker(self) -> None: + """W&B + TensorBoard tracker.""" + assert self.output_dir is not None and self.full_config is not None + self.tracker = ExperimentTracker( + output_dir=self.output_dir, + config=self.full_config, + use_wandb=self.tracking_cfg.use_wandb, + use_tb=self.tracking_cfg.use_tb, + wandb_project=self.tracking_cfg.wandb_project, + wandb_run_name=self.tracking_cfg.wandb_run_name, + wandb_entity=self.tracking_cfg.wandb_entity, + ) + + def _build_model(self) -> None: + """Build (or load) the encoder model based on the active backbone.""" + backbone = self.models_common_cfg.backbone + + if self.pipeline_cfg.resume_from is not None: + self._build_model_from_resume(backbone) + return + + # Fresh build. + mode_str = "baseline (no text)" if self.models_common_cfg.baseline_mode else "with text (L1/L2/L3)" + if backbone == "stripnet": + enc_str = "StripNet-small (shared, 512→1024 proj)" + else: # dinov3 + assert isinstance(self.models_cfg, DINOv3ModelsConfig) + enc_str = "shared DINOv3 WEB" if self.models_cfg.shared_encoder \ + else "asymmetric (WEB + SAT)" + LOGGER.info("Building model — %s, %s", mode_str, enc_str) + + if backbone == "stripnet": + self.model = self._build_stripnet_model() + else: # dinov3 + self.model = self._build_dinov3_model() + LOGGER.info("embed_dim=%d", self.model.embed_dim) + + def _build_model_from_resume(self, backbone: str) -> None: + """Resume model from checkpoint. Sets self.model, self.resume_ckpt, self.start_epoch.""" + LOGGER.info("Resuming from %s", self.pipeline_cfg.resume_from) + # Both DINOv3 and StripNet go through AsymmetricEncoder.load_checkpoint. + # Note: load_checkpoint doesn't support StripNet — known existing limitation. + if isinstance(self.models_cfg, DINOv3ModelsConfig): + dino_web_path = self.models_cfg.dino_web_path + dino_sat_path = self.models_cfg.dino_sat_path + else: + # StripNet preset on resume — fall back to original defaults. + dino_web_path = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth" + dino_sat_path = "nn_models/DINO_SAT/model.safetensors" + self.model, self.resume_ckpt = AsymmetricEncoder.load_checkpoint( + self.pipeline_cfg.resume_from, + dino_web_path=dino_web_path, + dino_sat_path=dino_sat_path, + lrsclip_path=self.models_common_cfg.lrsclip_path, + device=self.hardware_cfg.device, + ) + self.start_epoch = self.resume_ckpt.get("epoch", -1) + 1 + + def _build_stripnet_model(self) -> nn.Module: + """Construct AsymmetricEncoder configured for StripNet.""" + assert isinstance(self.models_cfg, StripNetModelsConfig) + m = self.models_cfg + # DINO paths passed but ignored at runtime when backbone='stripnet'. + # mona_bottleneck=64 matches the original TrainConfigGTAUAV default — + # used by inject_conv_mona_into_stripnet. + return AsymmetricEncoder( + dino_web_path="nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth", + dino_sat_path="nn_models/DINO_SAT/model.safetensors", + lrsclip_path=self.models_common_cfg.lrsclip_path, + init_gate=self.models_common_cfg.init_gate, + baseline_mode=self.models_common_cfg.baseline_mode, + shared_encoder=True, # StripNet always shared (overridden internally) + mona_bottleneck=64, # matches old TrainConfigGTAUAV.mona_bottleneck + mona_last_n_blocks=12, # not used for StripNet, but accepted by sig + device=self.hardware_cfg.device, + backbone="stripnet", + stripnet_path=m.stripnet_path, + stripnet_mona_last_n_stages=m.stripnet_mona_last_n_stages, + stripnet_freeze=m.stripnet_freeze, + ).to(self.hardware_cfg.device) + + def _build_dinov3_model(self) -> nn.Module: + """Construct AsymmetricEncoder configured for DINOv3.""" + assert isinstance(self.models_cfg, DINOv3ModelsConfig) + m = self.models_cfg + # stripnet_path passed with the original default — ignored at runtime + # for DINOv3 backbone. + return AsymmetricEncoder( + dino_web_path=m.dino_web_path, + dino_sat_path=m.dino_sat_path, + lrsclip_path=self.models_common_cfg.lrsclip_path, + init_gate=self.models_common_cfg.init_gate, + baseline_mode=self.models_common_cfg.baseline_mode, + shared_encoder=m.shared_encoder, + mona_bottleneck=m.mona_bottleneck, + mona_last_n_blocks=m.mona_last_n_blocks, + device=self.hardware_cfg.device, + backbone="dinov3", + stripnet_path="nn_models/STRIPNET/stripnet_s.pth", + stripnet_mona_last_n_stages=0, + stripnet_freeze=True, + ).to(self.hardware_cfg.device) + + def _configure_gradient_checkpointing(self) -> None: + """Enable gradient checkpointing on encoders that support it.""" + assert self.model is not None + backbone = self.models_common_cfg.backbone + if not self.hardware_cfg.gradient_checkpointing: + return + if backbone == "dinov3": + assert isinstance(self.models_cfg, DINOv3ModelsConfig) + if self.models_cfg.shared_encoder: + self.model.image_encoder.set_gradient_checkpointing(True) + else: + self.model.drone_encoder.set_gradient_checkpointing(True) + self.model.sat_encoder.set_gradient_checkpointing(True) + if self.model.text_encoder is not None: + self.model.text_encoder.transformer.gradient_checkpointing = True + LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)") + elif backbone == "stripnet": + if self.model.text_encoder is not None: + self.model.text_encoder.transformer.gradient_checkpointing = True + LOGGER.info("Gradient checkpointing enabled (DGTRS only; StripNet doesn't support it)") + + def _log_model_summary(self) -> None: + """Log trainable param count, save model_summary.txt, hook W&B.""" + assert self.model is not None and self.output_dir is not None and self.tracker is not None + n_trainable = sum(p.numel() for p in self.model.trainable_parameters()) + n_total = sum(p.numel() for p in self.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 = print_model_summary(self.model, device=self.hardware_cfg.device) + (self.output_dir / "model_summary.txt").write_text(model_summary) + if self.tracker.has_wandb: + self.tracker.watch_model(self.model, log_freq=50) + + def _build_loss(self) -> None: + """Build InfoNCELoss or WeightedInfoNCELoss based on training_cfg.loss_type.""" + t = self.training_cfg + if t.loss_type == "symmetric": + self.loss_fn = InfoNCELoss( + temperature_init=t.tau_init, + temperature_final=t.tau_final, + label_smoothing=t.label_smoothing, + weight_q2g=t.weight_q2g, + weight_g2q=t.weight_g2q, + learnable_temperature=t.learnable_temperature, + tau_min=t.tau_min, + tau_max=t.tau_max, + hard_mining_k=t.hard_mining_k, + ) + loss_name = "SymmetricInfoNCE" + elif t.loss_type == "weighted": + self.loss_fn = WeightedInfoNCELoss( + temperature_init=t.tau_init, + learnable_temperature=t.learnable_temperature, + label_smoothing=t.label_smoothing, + k=t.weighted_loss_k, + tau_min=t.tau_min, + tau_max=t.tau_max, + ) + loss_name = "WeightedInfoNCE" + else: + raise ValueError( + f"Unknown loss_type={t.loss_type!r} (expected 'symmetric' or 'weighted')", + ) + LOGGER.info( + "Loss: %s Temperature: %s (init=%.3f) q2g=%.2f g2q=%.2f", + loss_name, + "learnable" if t.learnable_temperature else "fixed", + t.tau_init, t.weight_q2g, t.weight_g2q, + ) + + def _build_neg_bank(self) -> None: + """Optional NegativeMemoryBank for hard-negative mining.""" + assert self.model is not None + if self.training_cfg.neg_bank_size > 0: + self.neg_bank = NegativeMemoryBank( + size=self.training_cfg.neg_bank_size, + dim=self.model.embed_dim, + ).to(self.hardware_cfg.device) + LOGGER.info( + "Negative memory bank: size=%d, dim=%d", + self.training_cfg.neg_bank_size, self.model.embed_dim, + ) + + def _build_data_loaders(self) -> None: + """Build train/test/train_eval datasets, samplers, loaders.""" + 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) + + self.train_ds = GTAUAVDataset( + pair_json=self.pipeline_cfg.train_json, + rgb_root=self.pipeline_cfg.rgb_root, + caption_root=self.pipeline_cfg.caption_root, + drone_transform=drone_train_tf, + sat_transform=sat_train_tf, + filter_meta=self.pipeline_cfg.filter_meta, + ) + self.test_ds = GTAUAVDataset( + pair_json=self.pipeline_cfg.test_json, + rgb_root=self.pipeline_cfg.rgb_root, + caption_root=self.pipeline_cfg.caption_root, + image_transform=eval_tf, + filter_meta=self.pipeline_cfg.filter_meta, + ) + self.train_eval_ds = GTAUAVDataset( + pair_json=self.pipeline_cfg.train_json, + rgb_root=self.pipeline_cfg.rgb_root, + caption_root=self.pipeline_cfg.caption_root, + image_transform=eval_tf, + filter_meta=self.pipeline_cfg.filter_meta, + ) + + self._build_batch_sampler() + + if self.batch_sampler is not None: + self.train_loader = DataLoader( + self.train_ds, + batch_sampler=self.batch_sampler, + num_workers=self.hardware_cfg.num_workers, + collate_fn=collate_gtauav_batch, + pin_memory=True, + ) + else: + self.train_loader = DataLoader( + self.train_ds, + batch_size=self.hardware_cfg.batch_size, + shuffle=True, + num_workers=self.hardware_cfg.num_workers, + collate_fn=collate_gtauav_batch, + pin_memory=True, + drop_last=True, + ) + + if self.training_cfg.dss_cache_dir is not None: + self.emb_cache = EmbeddingCache(self.training_cfg.dss_cache_dir) + LOGGER.info("DSS embedding cache: %s", self.training_cfg.dss_cache_dir) + + self.test_loader = DataLoader( + self.test_ds, + batch_size=self.hardware_cfg.batch_size, + shuffle=False, + num_workers=self.hardware_cfg.num_workers, + collate_fn=collate_gtauav_batch, + pin_memory=True, + ) + self.train_eval_loader = DataLoader( + self.train_eval_ds, + batch_size=self.hardware_cfg.batch_size, + shuffle=False, + num_workers=self.hardware_cfg.num_workers, + collate_fn=collate_gtauav_batch, + pin_memory=True, + ) + + effective_batch = self.hardware_cfg.batch_size * self.hardware_cfg.grad_accum_steps + LOGGER.info( + "train=%d test=%d batch=%d accum=%d effective_batch=%d", + len(self.train_ds), len(self.test_ds), + self.hardware_cfg.batch_size, self.hardware_cfg.grad_accum_steps, effective_batch, + ) + + def _build_batch_sampler(self) -> None: + """Choose between DSS / mutex / plain shuffle samplers.""" + assert self.train_ds is not None + sat_cand_list = [entry["sat_candidates"] for entry in self.train_ds.entries] + # Backward compat alias. + effective_sampler_type = ( + self.training_cfg.sampler_type + if self.training_cfg.use_mutex_sampler else "none" + ) + t = self.training_cfg + if effective_sampler_type == "dss": + self.batch_sampler = DynamicSimilaritySampler( + sat_cand_list, + batch_size=self.hardware_cfg.batch_size, + shuffle=True, + seed=self.pipeline_cfg.seed, + knn_device=t.dss_knn_device, + use_lsh=t.dss_use_lsh, + lsh_num_tables=t.dss_lsh_num_tables, + lsh_num_bits=t.dss_lsh_num_bits, + ) + LOGGER.info( + "Sampler: DynamicSimilarity — kNN on %s%s, warmup=%d, re-embed every %d epochs", + t.dss_knn_device, " + LSH" if t.dss_use_lsh else "", + t.dss_warmup_epochs, t.dss_reembed_every, + ) + elif effective_sampler_type == "mutex": + self.batch_sampler = MutuallyExclusiveSampler( + sat_cand_list, + batch_size=self.hardware_cfg.batch_size, + shuffle=True, + seed=self.pipeline_cfg.seed, + ) + LOGGER.info("Sampler: MutuallyExclusive — no false negatives within a batch") + else: + self.batch_sampler = None + LOGGER.info("Sampler: default shuffle (no mutex / no DSS)") + + def _build_optimizer_and_scheduler(self) -> None: + """Build AdamW with per-group LR + cosine-warmup scheduler + GradScaler.""" + assert self.model is not None and self.loss_fn is not None and self.train_loader is not None + t = self.training_cfg + + stripnet_lr_factor = ( + self.models_cfg.stripnet_backbone_lr_factor + if isinstance(self.models_cfg, StripNetModelsConfig) + else 0.1 + ) + param_groups = _build_param_groups( + self.model, + t.learning_rate, + t.text_lr_factor, + stripnet_backbone_lr_factor=stripnet_lr_factor, + ) + if t.learnable_temperature and self.loss_fn.logit_scale is not None: + param_groups[0]["params"].append(self.loss_fn.logit_scale) + self.optimizer = AdamW(param_groups, weight_decay=t.weight_decay) + + lr_info = f"proj={t.learning_rate:.0e}" + if not self.models_common_cfg.baseline_mode: + lr_info += f" text={t.learning_rate * t.text_lr_factor:.0e}" + LOGGER.info( + "Optimizer: AdamW LR: %s warmup=%d epochs", + lr_info, self.pipeline_cfg.warmup_epochs, + ) + + # Scheduler — cosine with linear warmup (counted in optimizer steps). + self.steps_per_epoch = math.ceil( + len(self.train_loader) / self.hardware_cfg.grad_accum_steps, + ) + total_steps = self.pipeline_cfg.epochs * self.steps_per_epoch + warmup_steps = self.pipeline_cfg.warmup_epochs * self.steps_per_epoch + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*") + self.scheduler = LambdaLR( + self.optimizer, + lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps), + last_epoch=-1, + ) + self.scaler = GradScaler(enabled=self.hardware_cfg.use_amp) + + def _restore_from_resume(self) -> None: + """Restore optimizer/scheduler/loss state on resume.""" + if self.resume_ckpt is None: + return + assert self.optimizer is not None and self.loss_fn is not None and self.scheduler is not None + if "optimizer_state" in self.resume_ckpt: + self.optimizer.load_state_dict(self.resume_ckpt["optimizer_state"]) + LOGGER.info("Optimizer state restored") + if "loss_state" in self.resume_ckpt: + self.loss_fn.load_state_dict(self.resume_ckpt["loss_state"]) + LOGGER.info("Loss state restored (tau=%.4f)", self.loss_fn.current_temperature) + # Set scheduler last_epoch so it resumes at the correct LR. + self.scheduler.last_epoch = self.start_epoch * self.steps_per_epoch + self.global_step = self.start_epoch * self.steps_per_epoch + LOGGER.info("Resuming from epoch %d", self.start_epoch) + + def _setup_profiler(self) -> None: + """Optional PyTorch profiler (only if start_epoch == 0).""" + if self.tracking_cfg.use_profiler and self.start_epoch == 0: + assert self.output_dir is not None + self.profiler = TrainingProfiler( + output_dir=self.output_dir, + n_warmup=self.tracking_cfg.profiler_warmup, + n_active=self.tracking_cfg.profiler_active, + ) + self.profiler.start() + + # =================================================================== + # Training loop + # =================================================================== + + def _train_loop(self) -> None: + """Main per-epoch loop.""" + LOGGER.info( + "Starting training for %d epochs (from epoch %d)", + self.pipeline_cfg.epochs, self.start_epoch, + ) + for epoch in range(self.start_epoch, self.pipeline_cfg.epochs): + self._maybe_dss_reembed(epoch) + epoch_start = time.time() + train_means = self._train_one_epoch(epoch) + elapsed = time.time() - epoch_start + self._log_epoch_summary(epoch, train_means, elapsed) + self._maybe_evaluate(epoch, train_means, elapsed) + self._save_checkpoint(epoch) + + def _maybe_dss_reembed(self, epoch: int) -> None: + """Refresh DSS query embeddings if the epoch is past warmup and on cadence.""" + if not isinstance(self.batch_sampler, DynamicSimilaritySampler): + return + t = self.training_cfg + if epoch < t.dss_warmup_epochs: + return + if (epoch - t.dss_warmup_epochs) % t.dss_reembed_every != 0: + return + + assert self.train_ds is not None and self.model is not None + query_embs: torch.Tensor | None = None + if self.emb_cache is not None: + query_embs = self.emb_cache.load(epoch) + if query_embs is None: + LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(self.train_ds), epoch) + t_embed = time.time() + query_embs = _embed_drone_queries( + self.model, self.train_ds, self.hardware_cfg.device, + batch_size=self.hardware_cfg.batch_size * self.hardware_cfg.grad_accum_steps, + num_workers=self.hardware_cfg.num_workers, + ) + LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed) + if self.emb_cache is not None: + self.emb_cache.save(epoch, query_embs) + t_sampler = time.time() + self.batch_sampler.update_embeddings(query_embs) + LOGGER.info("DSS: sampler update_embeddings took %.2fs", time.time() - t_sampler) + + def _train_one_epoch(self, epoch: int) -> dict[str, float]: + """One epoch of training. Returns mean metrics (loss, temperature, gates).""" + assert self.model is not None and self.loss_fn is not None + assert self.optimizer is not None and self.scheduler is not None and self.scaler is not None + assert self.train_loader is not None and self.tracker is not None and self.csv_logger is not None + + self.model.train() + if self.batch_sampler is not None: + self.batch_sampler.set_epoch(epoch) + + agg: dict[str, float] = {} + n_batches = 0 + accum = self.hardware_cfg.grad_accum_steps + device = self.hardware_cfg.device + baseline_mode = self.models_common_cfg.baseline_mode + + pbar = tqdm( + self.train_loader, + desc=f" Epoch {epoch + 1}/{self.pipeline_cfg.epochs}", + unit="batch", + leave=False, + ) + for batch in pbar: + if n_batches % accum == 0: + self.optimizer.zero_grad(set_to_none=True) + + drone_img = batch["drone_img"].to(device, non_blocking=True) + sat_img = batch["sat_img"].to(device, non_blocking=True) + altitude = batch.get("altitude") + if altitude is not None: + altitude = altitude.to(device, non_blocking=True) + + with autocast(device_type="cuda", enabled=self.hardware_cfg.use_amp): + if baseline_mode: + embeddings = self.model(drone_img=drone_img, sat_img=sat_img, altitude=altitude) + else: + embeddings = self.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, + ) + + queue_neg = self.neg_bank.get_queue() if self.neg_bank is not None else None + loss_kwargs: dict[str, Any] = { + "embeddings": embeddings, + "epoch": epoch, + "total_epochs": self.pipeline_cfg.epochs, + "queue_negatives": queue_neg, + } + if isinstance(self.loss_fn, WeightedInfoNCELoss): + loss_kwargs["positive_weights"] = batch["positive_weights"].to( + device, non_blocking=True, + ) + loss_dict = self.loss_fn(**loss_kwargs) + + raw_loss = float(loss_dict["total"].item()) + total_loss = loss_dict["total"] / accum + self.scaler.scale(total_loss).backward() + + if self.neg_bank is not None: + self.neg_bank.enqueue(embeddings["gallery"].detach()) + + is_accum_step = (n_batches + 1) % accum == 0 or (n_batches + 1) == len(self.train_loader) + if is_accum_step: + if self.training_cfg.grad_clip > 0: + self.scaler.unscale_(self.optimizer) + nn.utils.clip_grad_norm_( + self.model.trainable_parameters(), + max_norm=self.training_cfg.grad_clip, + ) + if self.tracking_cfg.log_grad_norms and n_batches % (50 * accum) < accum: + grad_norms = compute_gradient_norms(self.model, self.loss_fn) + self.tracker.log_gradients(epoch, grad_norms, step=self.global_step) + if n_batches < accum: + log_gradient_summary(grad_norms) + self.scaler.step(self.optimizer) + self.scaler.update() + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*") + self.scheduler.step() + self.global_step += 1 + + 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": self.optimizer.param_groups[0]["lr"], + } + self.tracker.log_train(epoch, step_metrics, step=self.global_step) + self.csv_logger.log_batch(epoch, n_batches, self.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}", + ) + + if self.profiler is not None: + self.profiler.step() + if self.profiler.is_done(n_batches): + self.profiler.export() + self.profiler = None + + return {k: v / max(n_batches, 1) for k, v in agg.items()} + + def _log_epoch_summary( + self, epoch: int, means: dict[str, float], elapsed: float, + ) -> None: + """Log epoch-level training summary + VRAM.""" + assert self.optimizer is not None and self.tracker is not None + LOGGER.info( + "epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate_q=%.4f gate_g=%.4f", + epoch, elapsed, + self.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), + ) + if torch.cuda.is_available(): + vram_gb = torch.cuda.max_memory_allocated() / 1e9 + self.tracker.log_scalar("system/vram_peak_gb", vram_gb, step=self.global_step) + + def _maybe_evaluate( + self, epoch: int, train_means: dict[str, float], elapsed: float, + ) -> None: + """Run eval if epoch % eval_every == 0 (or last epoch). Updates history + CSVs.""" + assert self.model is not None and self.loss_fn is not None + assert self.train_eval_loader is not None and self.test_loader is not None + assert self.csv_logger is not None and self.tracker is not None and self.optimizer is not None + + epoch_record: dict[str, Any] = { + "epoch": epoch, + "elapsed_seconds": elapsed, + "train": train_means, + } + + is_last = epoch == self.pipeline_cfg.epochs - 1 + run_eval = (epoch + 1) % self.pipeline_cfg.eval_every == 0 or is_last + + train_recall: dict[str, float] = {} + if run_eval: + train_eval_batches = len(self.test_loader) + train_recall = evaluate( + self.model, self.train_eval_loader, self.hardware_cfg.device, + loss_fn=self.loss_fn, epoch=epoch, total_epochs=self.pipeline_cfg.epochs, + max_batches=train_eval_batches, desc="eval-train", + ) + epoch_record["train_recall"] = train_recall + self.csv_logger.log_train_recall(epoch, train_recall) + self.tracker.log_train( + epoch, + {f"recall/{k}": v for k, v in train_recall.items() if k.startswith("r@")}, + step=self.global_step, + ) + + # Train CSV row (combines means + recall if eval ran). + train_row: dict[str, float] = {**train_means} + if "total" in train_row: + train_row["train_loss"] = train_row.pop("total") + if train_recall: + for key in ( + "r@1_q2g", "r@5_q2g", "r@10_q2g", "ap_q2g", + "r@1_g2q", "r@5_g2q", "r@10_g2q", "ap_g2q", + ): + train_row[key] = train_recall.get(key, 0.0) + self.csv_logger.log_train( + epoch, train_row, self.optimizer.param_groups[0]["lr"], elapsed, + ) + generate_plots(self.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_metrics = evaluate( + self.model, self.test_loader, self.hardware_cfg.device, + loss_fn=self.loss_fn, epoch=epoch, total_epochs=self.pipeline_cfg.epochs, + desc="eval-val", + ) + epoch_record["val"] = val_metrics + self.csv_logger.log_val(epoch, val_metrics) + generate_plots(self.csv_logger.log_dir) + self.tracker.log_val(epoch, val_metrics, step=self.global_step) + + r1 = val_metrics.get("r@1_q2g", 0.0) + if r1 > self.best_r1: + self.best_r1 = r1 + self.tracker.log_scalar( + "val/best_r@1_q2g", self.best_r1, step=self.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), + ) + + self._maybe_run_gradcam(epoch) + + self.history.append(epoch_record) + + def _maybe_run_gradcam(self, epoch: int) -> None: + """Optional Grad-CAM visualisation (only when gradcam_every divides epoch+1).""" + if not self.tracking_cfg.use_gradcam: + return + if (epoch + 1) % self.tracking_cfg.gradcam_every != 0: + return + from src.training.gradcam import generate_gradcam_samples + assert self.model is not None and self.test_loader is not None + assert self.output_dir is not None and self.tracker is not None + overlays = generate_gradcam_samples( + model=self.model, + dataloader=self.test_loader, + device=self.hardware_cfg.device, + output_dir=str(self.output_dir), + n_samples=self.tracking_cfg.gradcam_samples, + epoch=epoch, + ) + for i, overlay in enumerate(overlays[:4]): + kind = "drone" if i % 2 == 0 else "sat" + self.tracker.log_image( + f"gradcam/{kind}_{i // 2}", + overlay, + step=self.global_step, + caption=f"Epoch {epoch} {kind} Grad-CAM", + ) + + def _save_checkpoint(self, epoch: int) -> None: + """Save checkpoint atomically with backbone-specific arch flags.""" + assert self.model is not None and self.optimizer is not None + assert self.loss_fn is not None and self.output_dir is not None + backbone = self.models_common_cfg.backbone + ckpt_obj: dict[str, Any] = { + "epoch": epoch, + "model_state": self.model.state_dict(), + "optimizer_state": self.optimizer.state_dict(), + "loss_state": self.loss_fn.state_dict(), + "baseline_mode": self.models_common_cfg.baseline_mode, + "backbone": backbone, + } + if isinstance(self.models_cfg, DINOv3ModelsConfig): + ckpt_obj["shared_encoder"] = self.models_cfg.shared_encoder + ckpt_obj["mona_bottleneck"] = self.models_cfg.mona_bottleneck + ckpt_obj["mona_last_n_blocks"] = self.models_cfg.mona_last_n_blocks + # StripNet: no extra arch flags saved here (params come from gin on resume). + atomic_save_torch(ckpt_obj, self.output_dir / f"ckpt_epoch{epoch:03d}.pt") + LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch) + + # =================================================================== + # Final phase + # =================================================================== + + def _final_evaluation(self) -> None: + """Save history.json + eval_report.json + W&B summary.""" + assert self.model is not None and self.loss_fn is not None + assert self.test_loader is not None and self.output_dir is not None + assert self.tracker is not None and self.full_config is not None + + history_path = self.output_dir / "history.json" + with history_path.open("w", encoding="utf-8") as f: + json.dump(self.history, f, indent=2) + + LOGGER.info("Running final evaluation...") + final_metrics = evaluate( + self.model, self.test_loader, self.hardware_cfg.device, + loss_fn=self.loss_fn, + epoch=self.pipeline_cfg.epochs - 1, + total_epochs=self.pipeline_cfg.epochs, + ) + report = { + "config": self.full_config, + "metrics": final_metrics, + "history": self.history, + } + report_path = self.output_dir / "eval_report.json" + with report_path.open("w", encoding="utf-8") as f: + json.dump(report, f, indent=2) + + self.tracker.log_summary({ + "best_r@1_q2g": self.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), + }) + + 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 _cleanup(self) -> None: + """Close profiler + tracker.""" + if self.profiler is not None: + self.profiler.export() + self.profiler = None + if self.tracker is not None: + self.tracker.close() + + \ No newline at end of file