from __future__ import annotations """Training loop for caption quality test on cross-view geo-localization. GeoRSCLIP dual encoder with GatedFusion on query branch. Single InfoNCE loss: query(drone+text) vs gallery(satellite). """ import argparse import json import logging import time from pathlib import Path 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") @gin.configurable class TrainConfig: """Top-level training configuration. 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. """ def __init__( self, train_query_file: str = "Index/train_query.txt", val_query_file: str = "Index/test_query.txt", data_root: str = "/media/servml/SSD_2_TB/datasets/cvgl_datasets/UAV-GeoLoc", output_dir: str = "out/caption_test_exp_gate_SRGF", 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) if __name__ == "__main__": main()