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 `train()` plus dedicated `_setup_*` / `_build_*` / `_train_*` / `_evaluate_*` methods. Lifecycle: Trainer(...) → train() → done. `train()` 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: AsymmetricEncoder, 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 @torch.no_grad() def _embed_drone_queries( model: AsymmetricEncoder, train_ds: GTAUAVDataset, device: str, batch_size: int, num_workers: int, ) -> torch.Tensor: """Forward all drone queries and return [N, D] embeddings on CPU. Used by DynamicSimilaritySampler to rank drones by visual similarity. Runs with model.eval() but restores original train state afterwards. """ was_training = model.training model.eval() query_ds = GTAUAVDroneQuery(train_ds) loader = DataLoader( query_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_drone_query, pin_memory=True, ) embs: list[torch.Tensor] = [] for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False): drone_img = batch["drone_img"].to(device, non_blocking=True) q = model.encode_query( drone_img, batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], ) embs.append(q.cpu()) if was_training: model.train() return torch.cat(embs, dim=0) 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`. `train()` 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 train(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() #! ------ passed: OK -------------- 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.""" LOGGER.info("⚙️ Validate backbone") 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.""" LOGGER.info("⚙️ Setup out dir") 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.""" LOGGER.info("⚙️ Setup 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.""" LOGGER.info("⚙️ Build model...") 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("⚙️ Build model from resume...") 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.""" LOGGER.info("⚙️ Build StripNet model...") 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: LOGGER.info("⚙️ Build DINOv3 model...") """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.""" LOGGER.info("⚙️ Configure gradient checkpointing...") 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.""" LOGGER.info("⚙️ Build loss...") 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.""" LOGGER.info("⚙️ Build negative bank...") 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.""" LOGGER.info("⚙️ Build dataloaders...") 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.""" LOGGER.info("⚙️ Build optimizer & scheduler...") 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.""" LOGGER.info("⚙️ Restore from 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).""" LOGGER.info("⚙️ Setup profiler...") 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) 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) 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"] ) 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()