274 lines
8.2 KiB
Python
274 lines
8.2 KiB
Python
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()
|