Initial commit: caption quality test on UAV-VisLoc

Self-contained experimental track validating generated text captions
via retrieval R@1 lift on UAV-VisLoc.

Architecture: GeoRSCLIP ViT-B/32 dual encoder, 512-dim shared space.
Loss: 4-term InfoNCE (img-img + sat-cap + drone-cap + cap-cap)
      with cosine temperature decay, PALW-like curriculum.
Metric: delta R@1 (with text - without text) >= +3% => PASS.

Gin-configured (balanced / baseline_no_text / text_heavy variants).
Follows NADEZHDA code style.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

src/training/__init__.py

@@ -0,0 +1 @@
+"""Training loop for caption quality test."""

src/training/train.py

@@ -0,0 +1,313 @@
+from __future__ import annotations
+
+"""Training loop for caption quality validation on UAV-VisLoc.
+
+Uses gin-configurable DualEncoderCaptionTest + MultiTermInfoNCE.
+Logs per-component losses, temperature, and lambdas each step.
+Saves checkpoint + eval snapshot every epoch.
+"""
+
+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 (
+    VisLocCaptionDataset,
+    collate_caption_batch,
+)
+from src.eval.evaluate import evaluate_retrieval
+from src.losses.multi_infonce import MultiTermInfoNCE
+from src.models.dual_encoder import DualEncoderCaptionTest
+
+LOGGER = logging.getLogger("caption_test.train")
+
+
+@gin.configurable
+class TrainConfig:
+    """Top-level training configuration (gin-configurable).
+
+    Args:
+        train_manifest: Path to training manifest JSON.
+        val_manifest: Path to validation manifest JSON.
+        image_root: Directory prefix for images.
+        output_dir: Where to save checkpoints and logs.
+        epochs: Number of training epochs.
+        batch_size: Mini-batch size.
+        num_workers: DataLoader worker count.
+        learning_rate: AdamW initial LR.
+        weight_decay: AdamW weight decay.
+        grad_clip: Max gradient norm for clipping (0 disables).
+        use_amp: Enable fp16 mixed-precision training.
+        eval_every: Run validation every N epochs.
+        seed: Random seed.
+        device: torch device.
+    """
+
+    def __init__(
+        self,
+        train_manifest: str = "data/visloc_train.json",
+        val_manifest: str = "data/visloc_val.json",
+        image_root: str = "data/visloc/images",
+        output_dir: str = "out/caption_test",
+        epochs: int = 30,
+        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 = 1,
+        seed: int = 42,
+        device: str = "cuda",
+    ) -> None:
+        self.train_manifest = train_manifest
+        self.val_manifest = val_manifest
+        self.image_root = image_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:
+    """Seed Python, NumPy and PyTorch RNGs."""
+    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:
+    """Write torch checkpoint atomically (temp file + rename)."""
+    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 _step_loss(
+    model: DualEncoderCaptionTest,
+    loss_fn: MultiTermInfoNCE,
+    batch: dict,
+    epoch: int,
+    total_epochs: int,
+    device: str,
+    use_amp: bool,
+) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+    """Single training forward pass returning (total_loss, diagnostics)."""
+    drone_img = batch["drone_img"].to(device, non_blocking=True)
+    sat_img = batch["sat_img"].to(device, non_blocking=True)
+    caption_drone = batch["caption_drone"]
+    caption_sat = batch["caption_sat"]
+
+    with autocast(device_type="cuda", enabled=use_amp):
+        embeddings = model(
+            drone_img=drone_img,
+            sat_img=sat_img,
+            caption_drone=caption_drone,
+            caption_sat=caption_sat,
+        )
+        loss_dict = loss_fn(
+            embeddings=embeddings,
+            epoch=epoch,
+            total_epochs=total_epochs,
+        )
+
+    return loss_dict["total"], loss_dict
+
+
+def train(config_path: str) -> None:
+    """Run the full training loop driven by gin configuration.
+
+    Args:
+        config_path: Path to .gin config file.
+    """
+    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 = MultiTermInfoNCE().to(cfg.device)
+
+    # Datasets use the same preprocess function the model already holds.
+    preprocess = model.preprocess
+
+    train_ds = VisLocCaptionDataset(
+        manifest_path=cfg.train_manifest,
+        image_root=cfg.image_root,
+        image_transform=preprocess,
+    )
+    val_ds = VisLocCaptionDataset(
+        manifest_path=cfg.val_manifest,
+        image_root=cfg.image_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)
+
+    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)
+
+            total_loss, loss_dict = _step_loss(
+                model=model,
+                loss_fn=loss_fn,
+                batch=batch,
+                epoch=epoch,
+                total_epochs=cfg.epochs,
+                device=cfg.device,
+                use_amp=cfg.use_amp,
+            )
+
+            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()
+
+            # Accumulate diagnostics.
+            for key, tensor_val in loss_dict.items():
+                agg[key] = agg.get(key, 0.0) + float(tensor_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 total=%.4f img_img=%.4f "
+            "sat_cap=%.4f drone_cap=%.4f cap_cap=%.4f tau=%.4f",
+            epoch,
+            elapsed,
+            optimizer.param_groups[0]["lr"],
+            means.get("total", 0.0),
+            means.get("img_img", 0.0),
+            means.get("sat_cap", 0.0),
+            means.get("drone_cap", 0.0),
+            means.get("cap_cap", 0.0),
+            means.get("temperature", 0.0),
+        )
+
+        epoch_record = {
+            "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_d2s=%.4f R@1_s2d=%.4f "
+                "R@1_t2s=%.4f R@1_t2d=%.4f",
+                epoch,
+                val_metrics.get("r@1_drone_to_sat", 0.0),
+                val_metrics.get("r@1_sat_to_drone", 0.0),
+                val_metrics.get("r@1_text_to_sat", 0.0),
+                val_metrics.get("r@1_text_to_drone", 0.0),
+            )
+
+        history.append(epoch_record)
+
+        # Checkpoint per epoch.
+        _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",
+        )
+
+    # Save training history.
+    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="Path to gin configuration file.",
+    )
+    args = parser.parse_args()
+    train(config_path=args.config)
+
+
+if __name__ == "__main__":
+    main()