Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rewrite: GatedFusion architecture + UAV-GeoLoc dataset

Browse Source 
Architecture v2:
- Query branch: drone + text -> GatedFusion -> proj -> query_emb
- Gallery branch: satellite -> proj -> gallery_emb
- Single InfoNCE loss (asymmetric 0.6/0.4)
- GatedFusion: learnable gated addition (sigma(alpha)*img + (1-sigma(alpha))*text)
- Baseline mode: gate=1.0 (text ignored)

Dataset:
- UAV-GeoLoc loader with template captions from path metadata
- 27 terrain types with predefined features
- Random positive crop sampling per epoch

Configs: balanced (gate=0.7), baseline (no text), text_heavy (gate=0.3)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-17 17:13:00 +03:00
parent 2ce4017ebd
commit abb3337f8d
12 changed files with 1077 additions and 781 deletions
Download Patch File Download Diff File Expand all files Collapse all files

172
src/training/train.py
View File

@@ -1,10 +1,9 @@
from __future__ import annotations
"""Training loop for caption quality validation on UAV-VisLoc.
"""Training loop for caption quality test on cross-view geo-localization.
Uses gin-configurable DualEncoderCaptionTest + MultiTermInfoNCE.
Logs per-component losses, temperature, and lambdas each step.
Saves checkpoint + eval snapshot every epoch.
GeoRSCLIP dual encoder with GatedFusion on query branch.
Single InfoNCE loss: query(drone+text) vs gallery(satellite).
"""
import argparse
@@ -22,11 +21,11 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader
from src.datasets.visloc_with_captions import (
VisLocCaptionDataset,
GeoLocCaptionDataset,
collate_caption_batch,
)
from src.eval.evaluate import evaluate_retrieval
from src.losses.multi_infonce import MultiTermInfoNCE
from src.losses.multi_infonce import InfoNCELoss
from src.models.dual_encoder import DualEncoderCaptionTest
LOGGER = logging.getLogger("caption_test.train")
@@ -34,20 +33,20 @@ LOGGER = logging.getLogger("caption_test.train")
@gin.configurable
class TrainConfig:
"""Top-level training configuration (gin-configurable).
"""Top-level training configuration.
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.
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 worker count.
num_workers: DataLoader workers.
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.
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.
@@ -55,24 +54,24 @@ class TrainConfig:
def __init__(
self,
train_manifest: str = "data/visloc_train.json",
val_manifest: str = "data/visloc_val.json",
image_root: str = "data/visloc/images",
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 = 30,
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 = 1,
eval_every: int = 2,
seed: int = 42,
device: str = "cuda",
) -> None:
self.train_manifest = train_manifest
self.val_manifest = val_manifest
self.image_root = image_root
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
@@ -87,11 +86,8 @@ class TrainConfig:
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)
@@ -99,50 +95,14 @@ def _set_seed(seed: int) -> None:
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.
"""
"""Run full training loop from gin config."""
gin.parse_config_file(config_path)
cfg = TrainConfig()
@@ -153,21 +113,20 @@ def train(config_path: str) -> None:
_set_seed(cfg.seed)
cfg.output_dir.mkdir(parents=True, exist_ok=True)
# Model + loss
# Model + loss.
model = DualEncoderCaptionTest().to(cfg.device)
loss_fn = MultiTermInfoNCE().to(cfg.device)
loss_fn = InfoNCELoss().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,
train_ds = GeoLocCaptionDataset(
query_file=cfg.train_query_file,
data_root=cfg.data_root,
image_transform=preprocess,
)
val_ds = VisLocCaptionDataset(
manifest_path=cfg.val_manifest,
image_root=cfg.image_root,
val_ds = GeoLocCaptionDataset(
query_file=cfg.val_query_file,
data_root=cfg.data_root,
image_transform=preprocess,
)
@@ -197,6 +156,14 @@ def train(config_path: str) -> None:
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):
@@ -208,17 +175,25 @@ def train(config_path: str) -> None:
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,
)
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_(
@@ -228,9 +203,8 @@ def train(config_path: str) -> None:
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())
for key, val in loss_dict.items():
agg[key] = agg.get(key, 0.0) + float(val.item())
n_batches += 1
scheduler.step()
@@ -238,20 +212,15 @@ def train(config_path: str) -> None:
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,
"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("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),
means.get("gate", 1.0),
)
epoch_record = {
epoch_record: dict = {
"epoch": epoch,
"elapsed_seconds": elapsed,
"train": means,
@@ -267,18 +236,15 @@ def train(config_path: str) -> None:
)
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",
"val epoch=%d R@1_q2g=%.4f R@5_q2g=%.4f R@10_q2g=%.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),
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)
# Checkpoint per epoch.
_atomic_save(
obj={
"epoch": epoch,
@@ -289,7 +255,6 @@ def train(config_path: str) -> None:
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)
@@ -299,12 +264,7 @@ def train(config_path: str) -> None:
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.",
)
parser.add_argument("--config", type=str, required=True, help="Gin config file.")
args = parser.parse_args()
train(config_path=args.config)