1339 lines
55 KiB
Python
1339 lines
55 KiB
Python
from __future__ import annotations
|
||
|
||
"""Training loop for CVGL caption test on GTA-UAV-LR dataset.
|
||
|
||
Asymmetric DINOv3 encoders (drone LVD + satellite SAT) with LRSCLIP text fusion.
|
||
Single InfoNCE loss: query(drone+text) vs gallery(satellite).
|
||
|
||
Supports gin-config (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 <preset_name>
|
||
"""
|
||
|
||
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 <preset_name>",
|
||
)
|
||
|