From c43b4c82b9f193c43dff1df68a4f13ad9fffe006 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Mon, 4 May 2026 15:49:41 +0300 Subject: [PATCH] temp_step: verify actual final trainers scripts and remove obsolete --- src/eval/evaluator.py | 248 ++++++++++++++++++++++++++++++ src/training/csv_logger.py | 104 +++++++++++++ src/training/train.py | 307 ++++++------------------------------- 3 files changed, 400 insertions(+), 259 deletions(-) create mode 100644 src/eval/evaluator.py create mode 100644 src/training/csv_logger.py diff --git a/src/eval/evaluator.py b/src/eval/evaluator.py new file mode 100644 index 0000000..5ec1f3e --- /dev/null +++ b/src/eval/evaluator.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +"""Retrieval evaluation for GTA-UAV-LR cross-view geo-localization. + +Computes R@K and MRR for both q→g (drone→satellite) and g→q (satellite→drone) +on the full satellite gallery. Multi-match: a query counts as a hit@K if ANY +of its valid satellite matches (sat_candidates) appears in the top-K. + +Body transplanted from src/training/train_gtauav.py::_evaluate (pre-step-4b) +with two changes: + 1. Decorator @torch.no_grad() → @torch.inference_mode(). + 2. Type annotation `model: AsymmetricEncoder` → `model: nn.Module` + (any encoder with encode_query/encode_gallery + fusion_query.gate_value + and fusion_gallery.gate_value duck-typed attributes). + +Note: not to be confused with src/eval/evaluate.py (legacy v2 helper for +UAV-VisLoc with a different signature). This module lives at +src/eval/evaluator.py and is the active evaluator for v3 GTA-UAV-LR. +""" + +import logging +from typing import Any + +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +from tqdm import tqdm + +from src.datasets.gtauav_dataset import ( + GTAUAVDataset, + GTAUAVDroneQuery, + GTAUAVSatGallery, + collate_drone_query, + collate_sat_gallery, +) + +LOGGER = logging.getLogger("caption_test.evaluator") + + +@torch.inference_mode() +def evaluate( + model: nn.Module, + loader: DataLoader, + device: str, + loss_fn: nn.Module | None = None, + epoch: int = 0, + total_epochs: int = 1, + k_values: tuple[int, ...] = (1, 5, 10), + max_batches: int | None = None, + desc: str = "eval", +) -> dict[str, float]: + """Compute R@K and MRR on the full satellite gallery. + + Standard CVGL retrieval: forward every unique satellite in the dataset + once (gallery), forward every drone query, then rank gallery by + cosine similarity. A query counts as a hit@K if ANY of its valid + satellite matches (pair_pos_sate_img_list ∪ pair_pos_semipos_sate_img_list) + appears in the top-K. + + Args: + model: Encoder with `encode_query(drone_img, l1, l2, l3, altitude=...)` + and `encode_gallery(sat_img, l1, l2, l3)`. Must expose + `fusion_query.gate_value` and `fusion_gallery.gate_value`. + loader: DataLoader over a GTAUAVDataset (used only to pull dataset + + batch_size/num_workers/pin_memory; iteration is bypassed — + we build separate query and gallery loaders inside). + device: Torch device string. + loss_fn: If provided, computes per-batch loss against paired gallery + entries (uses the first valid sat per query as its positive). + The mean loss appears in the returned dict under 'loss'. + epoch, total_epochs: Passed through to loss_fn. + k_values: K values for R@K (e.g. (1, 5, 10)). + max_batches: Cap on query batches for quick sanity checks (gallery + is always full). + desc: tqdm description prefix. + + Returns: + Dict with: r@K_q2g, ap_q2g (= MRR), r@K_g2q, ap_g2q, loss (optional), + n_query, n_gallery, n_scored_g2q, gate_q, gate_g. + """ + 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 + + diff --git a/src/training/csv_logger.py b/src/training/csv_logger.py new file mode 100644 index 0000000..54b1f4a --- /dev/null +++ b/src/training/csv_logger.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +"""Per-batch and per-epoch CSV logger. + +Writes: + {output_dir}/logs/train.csv — epoch-level train averages + {output_dir}/logs/val.csv — epoch-level val metrics + {output_dir}/logs/train_recall.csv — epoch-level train recall metrics + {output_dir}/logs/train_batches.csv — per-batch train metrics (all epochs) + {output_dir}/logs/epoch_{N}_batches.csv — per-batch for one epoch + +Body transplanted verbatim from src/training/train_gtauav.py (pre-step-4b) +with no logic changes — only the relocation. +""" + +import logging +from pathlib import Path + +import pandas as pd + +LOGGER = logging.getLogger("caption_test.csv_logger") + + +class CSVLogger: + """Log train/val metrics to CSV files using pandas.""" + + 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, + ) + + \ No newline at end of file diff --git a/src/training/train.py b/src/training/train.py index 723d1c4..aecfac1 100644 --- a/src/training/train.py +++ b/src/training/train.py @@ -1,273 +1,62 @@ from __future__ import annotations -"""Training loop for caption quality test on cross-view geo-localization. +"""Thin wrapper around src.training.trainer.Trainer. -GeoRSCLIP dual encoder with GatedFusion on query branch. -Single InfoNCE loss: query(drone+text) vs gallery(satellite). +Kept for backward compatibility with src/main.py imports and any external +scripts that still call `train(...)` directly. After step 4b the body of +this module is just a delegation. + +Note: this module no longer runs standalone — entry point is src/main.py +(per REQUIREMENTS_GIN_STYLE.md §5): + python -m src.main """ -import argparse -import json -import logging -import time -from pathlib import Path +from src.conf.hardware_conf import HardwareConfig +from src.conf.models_common_conf import ModelsCommonConfig +from src.conf.models_dinov3_conf import DINOv3ModelsConfig +from src.conf.models_stripnet_conf import StripNetModelsConfig +from src.conf.pipeline_conf import PipelineConfig +from src.conf.tracking_conf import TrackingConfig +from src.conf.training_conf import TrainingConfig +from src.training.trainer_new import Trainer -import gin -import torch -import torch.nn as nn -from torch.amp import GradScaler, autocast -from torch.optim import AdamW -from torch.optim.lr_scheduler import CosineAnnealingLR -from torch.utils.data import DataLoader - -from src.datasets.visloc_with_captions import ( - GeoLocCaptionDataset, - collate_caption_batch, -) -from src.eval.evaluate import evaluate_retrieval -from src.losses.multi_infonce import InfoNCELoss -from src.models.dual_encoder import DualEncoderCaptionTest - -LOGGER = logging.getLogger("caption_test.train") +# Type alias re-exported for callers. +# SOFIA v1/v71 model configs exist in src/conf/ but are not yet supported +# by the trainer (no caption-aware fusion encoder wrapper). Adding them +# here will result in a NotImplementedError at Trainer.run(). +ModelsConfig = DINOv3ModelsConfig | StripNetModelsConfig -@gin.configurable -class TrainConfig: - """Top-level training configuration. +def train( + pipeline_cfg: PipelineConfig, + hardware_cfg: HardwareConfig, + training_cfg: TrainingConfig, + tracking_cfg: TrackingConfig, + models_common_cfg: ModelsCommonConfig, + models_cfg: ModelsConfig, +) -> None: + """Build a Trainer and run a full training cycle. Args: - train_query_file: Path to train_query.txt. - val_query_file: Path to test_query.txt (used as val). - data_root: Root of UAV-GeoLoc dataset. - output_dir: Checkpoint and log output directory. - epochs: Number of training epochs. - batch_size: Mini-batch size. - num_workers: DataLoader workers. - learning_rate: AdamW initial LR. - weight_decay: AdamW weight decay. - grad_clip: Max gradient norm (0 disables). - use_amp: Enable fp16 mixed-precision. - eval_every: Run validation every N epochs. - seed: Random seed. - device: torch device. + 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. """ - - def __init__( - self, - train_query_file: str = "Index/train_query.txt", - val_query_file: str = "Index/test_query.txt", - data_root: str = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc", - output_dir: str = "out/caption_test", - epochs: int = 10, - batch_size: int = 128, - num_workers: int = 4, - learning_rate: float = 1e-4, - weight_decay: float = 1e-4, - grad_clip: float = 1.0, - use_amp: bool = True, - eval_every: int = 2, - seed: int = 42, - device: str = "cuda", - ) -> None: - self.train_query_file = train_query_file - self.val_query_file = val_query_file - self.data_root = data_root - self.output_dir = Path(output_dir) - self.epochs = epochs - self.batch_size = batch_size - self.num_workers = num_workers - self.learning_rate = learning_rate - self.weight_decay = weight_decay - self.grad_clip = grad_clip - self.use_amp = use_amp - self.eval_every = eval_every - self.seed = seed - self.device = device - - -def _set_seed(seed: int) -> None: - import random as _random - import numpy as _np - _random.seed(seed) - _np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -def _atomic_save(obj: dict, path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - tmp_path = path.with_suffix(path.suffix + ".tmp") - torch.save(obj, tmp_path) - tmp_path.replace(path) - - -def train(config_path: str) -> None: - """Run full training loop from gin config.""" - gin.parse_config_file(config_path) - cfg = TrainConfig() - - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s %(name)s %(levelname)s %(message)s", - ) - _set_seed(cfg.seed) - cfg.output_dir.mkdir(parents=True, exist_ok=True) - - # Model + loss. - model = DualEncoderCaptionTest().to(cfg.device) - loss_fn = InfoNCELoss().to(cfg.device) - - preprocess = model.preprocess - - train_ds = GeoLocCaptionDataset( - query_file=cfg.train_query_file, - data_root=cfg.data_root, - image_transform=preprocess, - ) - val_ds = GeoLocCaptionDataset( - query_file=cfg.val_query_file, - data_root=cfg.data_root, - image_transform=preprocess, - ) - - train_loader = DataLoader( - train_ds, - batch_size=cfg.batch_size, - shuffle=True, - num_workers=cfg.num_workers, - collate_fn=collate_caption_batch, - pin_memory=True, - drop_last=True, - ) - val_loader = DataLoader( - val_ds, - batch_size=cfg.batch_size, - shuffle=False, - num_workers=cfg.num_workers, - collate_fn=collate_caption_batch, - pin_memory=True, - ) - - optimizer = AdamW( - model.trainable_parameters(), - lr=cfg.learning_rate, - weight_decay=cfg.weight_decay, - ) - scheduler = CosineAnnealingLR(optimizer, T_max=cfg.epochs) - scaler = GradScaler(enabled=cfg.use_amp) - - n_trainable = sum(p.numel() for p in model.trainable_parameters()) - n_total = sum(p.numel() for p in model.parameters()) - LOGGER.info( - "trainable=%d (%.2f%%) total=%d train=%d val=%d", - n_trainable, 100.0 * n_trainable / n_total, - n_total, len(train_ds), len(val_ds), - ) - - history: list[dict] = [] - - for epoch in range(cfg.epochs): - model.train() - epoch_start = time.time() - agg: dict[str, float] = {} - n_batches = 0 - - for batch in train_loader: - optimizer.zero_grad(set_to_none=True) - - drone_img = batch["drone_img"].to(cfg.device, non_blocking=True) - sat_img = batch["sat_img"].to(cfg.device, non_blocking=True) - caption_drone = batch["caption_drone"] - - with autocast(device_type="cuda", enabled=cfg.use_amp): - embeddings = model( - drone_img=drone_img, - sat_img=sat_img, - caption_drone=caption_drone, - ) - loss_dict = loss_fn( - embeddings=embeddings, - epoch=epoch, - total_epochs=cfg.epochs, - ) - - total_loss = loss_dict["total"] - scaler.scale(total_loss).backward() - - if cfg.grad_clip > 0: - scaler.unscale_(optimizer) - nn.utils.clip_grad_norm_( - model.trainable_parameters(), - max_norm=cfg.grad_clip, - ) - scaler.step(optimizer) - scaler.update() - - for key, val in loss_dict.items(): - agg[key] = agg.get(key, 0.0) + float(val.item()) - n_batches += 1 - - scheduler.step() - 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=%.4f", - epoch, elapsed, - optimizer.param_groups[0]["lr"], - means.get("total", 0.0), - means.get("temperature", 0.0), - means.get("gate", 1.0), - ) - - epoch_record: dict = { - "epoch": epoch, - "elapsed_seconds": elapsed, - "train": means, - } - - # Validation. - if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1: - model.eval() - val_metrics = evaluate_retrieval( - model=model, - loader=val_loader, - device=cfg.device, - ) - epoch_record["val"] = val_metrics - LOGGER.info( - "val epoch=%d R@1_q2g=%.4f R@5_q2g=%.4f R@10_q2g=%.4f", - epoch, - val_metrics.get("r@1_query_to_gallery", 0.0), - val_metrics.get("r@5_query_to_gallery", 0.0), - val_metrics.get("r@10_query_to_gallery", 0.0), - ) - - history.append(epoch_record) - - _atomic_save( - obj={ - "epoch": epoch, - "model_state": model.state_dict(), - "optimizer_state": optimizer.state_dict(), - "config_path": config_path, - }, - path=cfg.output_dir / f"ckpt_epoch{epoch:03d}.pt", - ) - - history_path = cfg.output_dir / "history.json" - with history_path.open("w", encoding="utf-8") as f: - json.dump(history, f, indent=2) - - LOGGER.info("training complete, history saved to %s", history_path) - - -def main() -> None: - parser = argparse.ArgumentParser(description="Caption quality test training.") - parser.add_argument("--config", type=str, required=True, help="Gin config file.") - args = parser.parse_args() - train(config_path=args.config) + Trainer( + pipeline_cfg=pipeline_cfg, + hardware_cfg=hardware_cfg, + training_cfg=training_cfg, + tracking_cfg=tracking_cfg, + models_common_cfg=models_common_cfg, + models_cfg=models_cfg, + ).run() if __name__ == "__main__": - main() + raise SystemExit( + "Direct execution removed. Use: python -m src.main ", + ) +