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 ", )