temp_step: verify actual final trainers scripts and remove obsolete
This commit is contained in:
pikaliov
248
src/eval/evaluator.py
Normal file
248
src/eval/evaluator.py
Normal file
@@ -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
|
||||
|
||||
|
||||
104
src/training/csv_logger.py
Normal file
104
src/training/csv_logger.py
Normal file
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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 <preset_name>
|
||||
"""
|
||||
|
||||
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 <preset_name>",
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user