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>
This commit is contained in:
1
src/__init__.py
Normal file
1
src/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Caption quality test package."""
|
||||
7
src/datasets/__init__.py
Normal file
7
src/datasets/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
"""Dataset loaders for caption quality test."""
|
||||
from src.datasets.visloc_with_captions import (
|
||||
VisLocCaptionDataset,
|
||||
collate_caption_batch,
|
||||
)
|
||||
|
||||
__all__ = ["VisLocCaptionDataset", "collate_caption_batch"]
|
||||
158
src/datasets/visloc_with_captions.py
Normal file
158
src/datasets/visloc_with_captions.py
Normal file
@@ -0,0 +1,158 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""UAV-VisLoc dataset loader augmented with generated captions.
|
||||
|
||||
Expects a manifest JSON of the form:
|
||||
[
|
||||
{
|
||||
"pair_id": "v001_0042",
|
||||
"drone_path": "drone/v001_0042.jpg",
|
||||
"sat_path": "satellite/v001_0042.png",
|
||||
"caption_drone": "low-altitude photo of residential ...",
|
||||
"caption_sat": "aerial view of urban area ...",
|
||||
"gps": [lat, lon]
|
||||
},
|
||||
...
|
||||
]
|
||||
|
||||
Captions are produced offline by scripts/generate_captions.py using one of
|
||||
three strategies: template, VLM, or hybrid (see АНАЛИЗ_caption_quality_test).
|
||||
"""
|
||||
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
import gin
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class VisLocCaptionDataset(Dataset):
|
||||
"""UAV-VisLoc pairs with generated captions.
|
||||
|
||||
Args:
|
||||
manifest_path: Path to JSON manifest with pair entries.
|
||||
image_root: Directory prefix joined with manifest relative paths.
|
||||
image_transform: Callable applied to PIL images (e.g., GeoRSCLIP preprocess).
|
||||
caption_strategy: Which caption field to use ('template', 'vlm', 'hybrid').
|
||||
The corresponding field must exist in the manifest
|
||||
(e.g., 'caption_sat_vlm', or the generic 'caption_sat').
|
||||
drop_caption_prob: Random probability of replacing a caption with ''.
|
||||
Useful for dropout ablations during training.
|
||||
seed: Random seed for reproducibility.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
manifest_path: str,
|
||||
image_root: str,
|
||||
image_transform: Callable[[Image.Image], torch.Tensor],
|
||||
caption_strategy: str = "hybrid",
|
||||
drop_caption_prob: float = 0.0,
|
||||
seed: int = 0,
|
||||
) -> None:
|
||||
self.manifest_path = Path(manifest_path)
|
||||
self.image_root = Path(image_root)
|
||||
self.image_transform = image_transform
|
||||
self.caption_strategy = caption_strategy
|
||||
self.drop_caption_prob = drop_caption_prob
|
||||
self._rng = random.Random(seed)
|
||||
|
||||
with self.manifest_path.open("r", encoding="utf-8") as f:
|
||||
self.entries: list[dict[str, Any]] = json.load(f)
|
||||
|
||||
self._validate_entries()
|
||||
|
||||
def _validate_entries(self) -> None:
|
||||
"""Ensure all entries have required fields for the chosen strategy."""
|
||||
required = {"drone_path", "sat_path"}
|
||||
caption_sat_key = self._caption_key("sat")
|
||||
caption_drone_key = self._caption_key("drone")
|
||||
required |= {caption_sat_key, caption_drone_key}
|
||||
|
||||
for i, entry in enumerate(self.entries):
|
||||
missing = required - entry.keys()
|
||||
if missing:
|
||||
raise KeyError(
|
||||
f"Entry {i} (pair_id={entry.get('pair_id', '?')}) missing fields: "
|
||||
f"{sorted(missing)}"
|
||||
)
|
||||
|
||||
def _caption_key(self, view: str) -> str:
|
||||
"""Resolve caption field name from strategy + view."""
|
||||
if self.caption_strategy == "hybrid":
|
||||
return f"caption_{view}"
|
||||
return f"caption_{view}_{self.caption_strategy}"
|
||||
|
||||
def _load_image(self, relative_path: str) -> torch.Tensor:
|
||||
"""Load image and apply preprocessing."""
|
||||
path = self.image_root / relative_path
|
||||
with Image.open(path) as img:
|
||||
rgb = img.convert("RGB")
|
||||
return self.image_transform(rgb)
|
||||
|
||||
def _maybe_drop(self, caption: str) -> str:
|
||||
"""Stochastically drop caption to empty string for robustness training."""
|
||||
if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
|
||||
return ""
|
||||
return caption
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.entries)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict[str, Any]:
|
||||
"""Return one pair with images and captions.
|
||||
|
||||
Args:
|
||||
idx: Index into the manifest.
|
||||
|
||||
Returns:
|
||||
Dict with:
|
||||
- 'drone_img': [3, H, W] tensor
|
||||
- 'sat_img': [3, H, W] tensor
|
||||
- 'caption_drone': str (possibly empty)
|
||||
- 'caption_sat': str (possibly empty)
|
||||
- 'pair_id': str for logging
|
||||
"""
|
||||
entry = self.entries[idx]
|
||||
|
||||
drone_img = self._load_image(entry["drone_path"])
|
||||
sat_img = self._load_image(entry["sat_path"])
|
||||
|
||||
caption_drone = self._maybe_drop(entry[self._caption_key("drone")])
|
||||
caption_sat = self._maybe_drop(entry[self._caption_key("sat")])
|
||||
|
||||
return {
|
||||
"drone_img": drone_img,
|
||||
"sat_img": sat_img,
|
||||
"caption_drone": caption_drone,
|
||||
"caption_sat": caption_sat,
|
||||
"pair_id": entry.get("pair_id", f"idx_{idx}"),
|
||||
}
|
||||
|
||||
|
||||
def collate_caption_batch(
|
||||
batch: list[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""Collate VisLocCaptionDataset items into a batched dict.
|
||||
|
||||
Images are stacked; captions remain Python lists so the tokenizer can
|
||||
process them inside the model.forward().
|
||||
|
||||
Args:
|
||||
batch: List of samples from VisLocCaptionDataset.__getitem__.
|
||||
|
||||
Returns:
|
||||
Batched dict with stacked image tensors and caption lists.
|
||||
"""
|
||||
return {
|
||||
"drone_img": torch.stack([b["drone_img"] for b in batch], dim=0),
|
||||
"sat_img": torch.stack([b["sat_img"] for b in batch], dim=0),
|
||||
"caption_drone": [b["caption_drone"] for b in batch],
|
||||
"caption_sat": [b["caption_sat"] for b in batch],
|
||||
"pair_ids": [b["pair_id"] for b in batch],
|
||||
}
|
||||
8
src/eval/__init__.py
Normal file
8
src/eval/__init__.py
Normal file
@@ -0,0 +1,8 @@
|
||||
"""Evaluation utilities for caption quality test."""
|
||||
from src.eval.evaluate import (
|
||||
delta_r_at_1,
|
||||
evaluate_retrieval,
|
||||
run_evaluation_from_checkpoint,
|
||||
)
|
||||
|
||||
__all__ = ["delta_r_at_1", "evaluate_retrieval", "run_evaluation_from_checkpoint"]
|
||||
242
src/eval/evaluate.py
Normal file
242
src/eval/evaluate.py
Normal file
@@ -0,0 +1,242 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""Evaluation utilities for caption quality test.
|
||||
|
||||
Implements retrieval metrics across four directions and a
|
||||
`delta_r_at_1` helper that compares caption-aware vs. image-only runs.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import gin
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from src.models.dual_encoder import DualEncoderCaptionTest
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.eval")
|
||||
|
||||
|
||||
def _recall_at_k(
|
||||
similarity: torch.Tensor,
|
||||
k_values: tuple[int, ...] = (1, 5, 10),
|
||||
) -> dict[int, float]:
|
||||
"""Compute Recall@K assuming positives on the diagonal.
|
||||
|
||||
Args:
|
||||
similarity: Pairwise similarity matrix [N_query, N_gallery].
|
||||
k_values: Tuple of K values to compute.
|
||||
|
||||
Returns:
|
||||
Dict mapping K -> recall in [0, 1].
|
||||
"""
|
||||
n_query = similarity.size(0)
|
||||
targets = torch.arange(n_query, device=similarity.device)
|
||||
sorted_idx = similarity.argsort(dim=1, descending=True)
|
||||
|
||||
result: dict[int, float] = {}
|
||||
for k in k_values:
|
||||
top_k = sorted_idx[:, :k]
|
||||
hit = (top_k == targets.unsqueeze(1)).any(dim=1).float()
|
||||
result[k] = float(hit.mean().item())
|
||||
return result
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def _encode_dataset(
|
||||
model: DualEncoderCaptionTest,
|
||||
loader: DataLoader,
|
||||
device: str,
|
||||
include_captions: bool,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Encode every sample in the loader into the shared embedding space.
|
||||
|
||||
Args:
|
||||
model: Trained dual encoder.
|
||||
loader: DataLoader yielding collated batches.
|
||||
device: Target device string.
|
||||
include_captions: If False, caption embeddings are skipped.
|
||||
|
||||
Returns:
|
||||
Dict with keys 'drone', 'sat', 'cap_drone', 'cap_sat' -> [N, D].
|
||||
"""
|
||||
model.eval()
|
||||
all_drone: list[torch.Tensor] = []
|
||||
all_sat: list[torch.Tensor] = []
|
||||
all_cap_drone: list[torch.Tensor] = []
|
||||
all_cap_sat: list[torch.Tensor] = []
|
||||
|
||||
for batch in loader:
|
||||
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
||||
sat_img = batch["sat_img"].to(device, non_blocking=True)
|
||||
captions_drone = batch["caption_drone"] if include_captions else None
|
||||
captions_sat = batch["caption_sat"] if include_captions else None
|
||||
|
||||
embeddings = model(
|
||||
drone_img=drone_img,
|
||||
sat_img=sat_img,
|
||||
caption_drone=captions_drone,
|
||||
caption_sat=captions_sat,
|
||||
)
|
||||
all_drone.append(embeddings["drone"].cpu())
|
||||
all_sat.append(embeddings["sat"].cpu())
|
||||
if include_captions:
|
||||
all_cap_drone.append(embeddings["cap_drone"].cpu())
|
||||
all_cap_sat.append(embeddings["cap_sat"].cpu())
|
||||
|
||||
out = {
|
||||
"drone": torch.cat(all_drone, dim=0),
|
||||
"sat": torch.cat(all_sat, dim=0),
|
||||
}
|
||||
if include_captions:
|
||||
out["cap_drone"] = torch.cat(all_cap_drone, dim=0)
|
||||
out["cap_sat"] = torch.cat(all_cap_sat, dim=0)
|
||||
return out
|
||||
|
||||
|
||||
def evaluate_retrieval(
|
||||
model: DualEncoderCaptionTest,
|
||||
loader: DataLoader,
|
||||
device: str,
|
||||
k_values: tuple[int, ...] = (1, 5, 10),
|
||||
include_captions: bool = True,
|
||||
) -> dict[str, float]:
|
||||
"""Compute retrieval metrics across four directions.
|
||||
|
||||
Directions reported (when captions included):
|
||||
drone -> sat, sat -> drone, text -> sat, text -> drone.
|
||||
|
||||
Args:
|
||||
model: Trained DualEncoderCaptionTest.
|
||||
loader: DataLoader over evaluation split.
|
||||
device: torch device string.
|
||||
k_values: Recall@K cutoffs.
|
||||
include_captions: If False, only image-image directions computed.
|
||||
|
||||
Returns:
|
||||
Flat dict with keys like 'r@1_drone_to_sat', 'r@5_text_to_sat', etc.
|
||||
"""
|
||||
feats = _encode_dataset(
|
||||
model=model,
|
||||
loader=loader,
|
||||
device=device,
|
||||
include_captions=include_captions,
|
||||
)
|
||||
|
||||
metrics: dict[str, float] = {}
|
||||
|
||||
sim_d2s = feats["drone"] @ feats["sat"].t()
|
||||
sim_s2d = sim_d2s.t()
|
||||
|
||||
for k, val in _recall_at_k(sim_d2s, k_values).items():
|
||||
metrics[f"r@{k}_drone_to_sat"] = val
|
||||
for k, val in _recall_at_k(sim_s2d, k_values).items():
|
||||
metrics[f"r@{k}_sat_to_drone"] = val
|
||||
|
||||
if include_captions and "cap_sat" in feats and "cap_drone" in feats:
|
||||
sim_t2s = feats["cap_sat"] @ feats["sat"].t()
|
||||
sim_t2d = feats["cap_drone"] @ feats["drone"].t()
|
||||
sim_tcd2tcs = feats["cap_drone"] @ feats["cap_sat"].t()
|
||||
|
||||
for k, val in _recall_at_k(sim_t2s, k_values).items():
|
||||
metrics[f"r@{k}_text_to_sat"] = val
|
||||
for k, val in _recall_at_k(sim_t2d, k_values).items():
|
||||
metrics[f"r@{k}_text_to_drone"] = val
|
||||
for k, val in _recall_at_k(sim_tcd2tcs, k_values).items():
|
||||
metrics[f"r@{k}_capdrone_to_capsat"] = val
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
def delta_r_at_1(
|
||||
full_metrics: dict[str, float],
|
||||
baseline_metrics: dict[str, float],
|
||||
direction: str = "drone_to_sat",
|
||||
) -> float:
|
||||
"""Compute caption-quality proxy: R@1 gain from adding captions.
|
||||
|
||||
Args:
|
||||
full_metrics: Metrics from training WITH caption losses.
|
||||
baseline_metrics: Metrics from training WITHOUT caption losses.
|
||||
direction: Retrieval direction to compare.
|
||||
|
||||
Returns:
|
||||
Δ R@1 in [−1, +1] range (positive = captions help).
|
||||
"""
|
||||
key = f"r@1_{direction}"
|
||||
if key not in full_metrics or key not in baseline_metrics:
|
||||
raise KeyError(
|
||||
f"Missing '{key}' in one of the metric dicts. "
|
||||
f"Available full={list(full_metrics)}, baseline={list(baseline_metrics)}"
|
||||
)
|
||||
return full_metrics[key] - baseline_metrics[key]
|
||||
|
||||
|
||||
@gin.configurable
|
||||
def run_evaluation_from_checkpoint(
|
||||
checkpoint_path: str,
|
||||
test_manifest: str,
|
||||
image_root: str,
|
||||
output_path: str = "eval_report.json",
|
||||
batch_size: int = 128,
|
||||
num_workers: int = 4,
|
||||
device: str = "cuda",
|
||||
) -> dict[str, float]:
|
||||
"""Standalone evaluation entry point (gin-configurable).
|
||||
|
||||
Args:
|
||||
checkpoint_path: Path to .pt checkpoint from training.
|
||||
test_manifest: Path to test manifest JSON.
|
||||
image_root: Directory prefix for images.
|
||||
output_path: Where to write the JSON report.
|
||||
batch_size: Batch size for encoding.
|
||||
num_workers: DataLoader workers.
|
||||
device: torch device.
|
||||
|
||||
Returns:
|
||||
Dict of retrieval metrics.
|
||||
"""
|
||||
from src.datasets.visloc_with_captions import (
|
||||
VisLocCaptionDataset,
|
||||
collate_caption_batch,
|
||||
)
|
||||
|
||||
model = DualEncoderCaptionTest().to(device)
|
||||
ckpt = torch.load(checkpoint_path, map_location=device)
|
||||
model.load_state_dict(ckpt["model_state"])
|
||||
model.eval()
|
||||
|
||||
test_ds = VisLocCaptionDataset(
|
||||
manifest_path=test_manifest,
|
||||
image_root=image_root,
|
||||
image_transform=model.preprocess,
|
||||
)
|
||||
test_loader = DataLoader(
|
||||
test_ds,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=num_workers,
|
||||
collate_fn=collate_caption_batch,
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
metrics = evaluate_retrieval(
|
||||
model=model,
|
||||
loader=test_loader,
|
||||
device=device,
|
||||
)
|
||||
|
||||
report = {
|
||||
"checkpoint": checkpoint_path,
|
||||
"test_manifest": test_manifest,
|
||||
"metrics": metrics,
|
||||
}
|
||||
out = Path(output_path)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
with out.open("w", encoding="utf-8") as f:
|
||||
json.dump(report, f, indent=2)
|
||||
|
||||
LOGGER.info("evaluation report saved to %s", out)
|
||||
return metrics
|
||||
8
src/losses/__init__.py
Normal file
8
src/losses/__init__.py
Normal file
@@ -0,0 +1,8 @@
|
||||
"""Loss functions for caption quality test."""
|
||||
from src.losses.multi_infonce import (
|
||||
MultiTermInfoNCE,
|
||||
cosine_temperature,
|
||||
curriculum_lambdas,
|
||||
)
|
||||
|
||||
__all__ = ["MultiTermInfoNCE", "cosine_temperature", "curriculum_lambdas"]
|
||||
262
src/losses/multi_infonce.py
Normal file
262
src/losses/multi_infonce.py
Normal file
@@ -0,0 +1,262 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""Multi-term InfoNCE loss for caption quality validation.
|
||||
|
||||
Four InfoNCE terms over projected embeddings:
|
||||
L = lambda_ii * L_img_img
|
||||
+ lambda_sc * L_sat_cap
|
||||
+ lambda_dc * L_drone_cap
|
||||
+ lambda_cc * L_cap_cap
|
||||
where L_img_img is the classical symmetric CVGL contrastive loss
|
||||
with asymmetric weights (0.6 drone->sat + 0.4 sat->drone).
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import gin
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def _symmetric_info_nce(
|
||||
emb_a: torch.Tensor,
|
||||
emb_b: torch.Tensor,
|
||||
temperature: float,
|
||||
label_smoothing: float,
|
||||
weight_a2b: float = 0.5,
|
||||
weight_b2a: float = 0.5,
|
||||
) -> torch.Tensor:
|
||||
"""Compute weighted symmetric InfoNCE between two L2-normalized embeddings.
|
||||
|
||||
Args:
|
||||
emb_a: First embedding set [B, D].
|
||||
emb_b: Second embedding set [B, D]. Positive pairs are on the diagonal.
|
||||
temperature: Softmax temperature (smaller = sharper distribution).
|
||||
label_smoothing: Cross-entropy label smoothing epsilon.
|
||||
weight_a2b: Weight for A-query direction.
|
||||
weight_b2a: Weight for B-query direction.
|
||||
|
||||
Returns:
|
||||
Scalar weighted loss.
|
||||
"""
|
||||
batch_size = emb_a.size(0)
|
||||
logits = emb_a @ emb_b.t() / temperature
|
||||
targets = torch.arange(batch_size, device=emb_a.device)
|
||||
|
||||
loss_a2b = F.cross_entropy(logits, targets, label_smoothing=label_smoothing)
|
||||
loss_b2a = F.cross_entropy(logits.t(), targets, label_smoothing=label_smoothing)
|
||||
|
||||
return weight_a2b * loss_a2b + weight_b2a * loss_b2a
|
||||
|
||||
|
||||
def cosine_temperature(
|
||||
epoch: int,
|
||||
total_epochs: int,
|
||||
tau_init: float = 0.1,
|
||||
tau_final: float = 0.01,
|
||||
) -> float:
|
||||
"""Cosine-decay schedule for InfoNCE temperature.
|
||||
|
||||
Args:
|
||||
epoch: Current training epoch (0-indexed).
|
||||
total_epochs: Total number of epochs.
|
||||
tau_init: Initial temperature.
|
||||
tau_final: Final temperature.
|
||||
|
||||
Returns:
|
||||
Temperature value for this epoch.
|
||||
"""
|
||||
total_epochs = max(total_epochs, 1)
|
||||
progress = min(max(epoch / total_epochs, 0.0), 1.0)
|
||||
cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
|
||||
return tau_final + (tau_init - tau_final) * cosine
|
||||
|
||||
|
||||
def curriculum_lambdas(
|
||||
epoch: int,
|
||||
warmup_epochs: int = 3,
|
||||
text_ramp_epochs: int = 10,
|
||||
lambda_ii: float = 1.0,
|
||||
lambda_sc_max: float = 0.3,
|
||||
lambda_dc_max: float = 0.3,
|
||||
lambda_cc_max: float = 0.1,
|
||||
) -> dict[str, float]:
|
||||
"""Compute per-epoch loss weights under the curriculum schedule.
|
||||
|
||||
- Epochs 0..warmup_epochs: image-image only.
|
||||
- Epochs warmup..text_ramp_epochs: linearly ramp sat-cap and drone-cap.
|
||||
- Epochs >= text_ramp_epochs: full loss including caption-caption term.
|
||||
|
||||
Args:
|
||||
epoch: Current epoch (0-indexed).
|
||||
warmup_epochs: Number of warmup epochs (no text losses).
|
||||
text_ramp_epochs: Epoch when text losses reach max.
|
||||
lambda_ii: Constant weight for image-image loss.
|
||||
lambda_sc_max: Max weight for satellite-caption loss.
|
||||
lambda_dc_max: Max weight for drone-caption loss.
|
||||
lambda_cc_max: Max weight for caption-caption loss.
|
||||
|
||||
Returns:
|
||||
Dict with keys 'img_img', 'sat_cap', 'drone_cap', 'cap_cap'.
|
||||
"""
|
||||
if epoch < warmup_epochs:
|
||||
ramp = 0.0
|
||||
elif epoch >= text_ramp_epochs:
|
||||
ramp = 1.0
|
||||
else:
|
||||
denom = max(text_ramp_epochs - warmup_epochs, 1)
|
||||
ramp = (epoch - warmup_epochs) / denom
|
||||
|
||||
return {
|
||||
"img_img": lambda_ii,
|
||||
"sat_cap": lambda_sc_max * ramp,
|
||||
"drone_cap": lambda_dc_max * ramp,
|
||||
"cap_cap": lambda_cc_max * ramp,
|
||||
}
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class MultiTermInfoNCE(nn.Module):
|
||||
"""Multi-term InfoNCE loss with curriculum and cosine temperature.
|
||||
|
||||
Produces total loss and per-component diagnostics. All inputs must be
|
||||
L2-normalized embeddings of the same dimension.
|
||||
|
||||
Args:
|
||||
temperature_init: Initial temperature (epoch 0).
|
||||
temperature_final: Final temperature after cosine decay.
|
||||
label_smoothing: Cross-entropy label smoothing epsilon.
|
||||
asym_drone_to_sat: Weight for drone->sat InfoNCE direction.
|
||||
asym_sat_to_drone: Weight for sat->drone InfoNCE direction.
|
||||
warmup_epochs: Epochs with image-image loss only.
|
||||
text_ramp_epochs: Epoch at which text losses reach max.
|
||||
lambda_ii: Constant weight for image-image loss.
|
||||
lambda_sc_max: Max weight for sat-caption loss.
|
||||
lambda_dc_max: Max weight for drone-caption loss.
|
||||
lambda_cc_max: Max weight for caption-caption loss.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
temperature_init: float = 0.1,
|
||||
temperature_final: float = 0.01,
|
||||
label_smoothing: float = 0.1,
|
||||
asym_drone_to_sat: float = 0.6,
|
||||
asym_sat_to_drone: float = 0.4,
|
||||
warmup_epochs: int = 3,
|
||||
text_ramp_epochs: int = 10,
|
||||
lambda_ii: float = 1.0,
|
||||
lambda_sc_max: float = 0.3,
|
||||
lambda_dc_max: float = 0.3,
|
||||
lambda_cc_max: float = 0.1,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.temperature_init = temperature_init
|
||||
self.temperature_final = temperature_final
|
||||
self.label_smoothing = label_smoothing
|
||||
self.asym_drone_to_sat = asym_drone_to_sat
|
||||
self.asym_sat_to_drone = asym_sat_to_drone
|
||||
self.warmup_epochs = warmup_epochs
|
||||
self.text_ramp_epochs = text_ramp_epochs
|
||||
self.lambda_ii = lambda_ii
|
||||
self.lambda_sc_max = lambda_sc_max
|
||||
self.lambda_dc_max = lambda_dc_max
|
||||
self.lambda_cc_max = lambda_cc_max
|
||||
|
||||
def forward(
|
||||
self,
|
||||
embeddings: dict[str, torch.Tensor],
|
||||
epoch: int,
|
||||
total_epochs: int,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Compute multi-term loss.
|
||||
|
||||
Args:
|
||||
embeddings: Dict with keys 'drone', 'sat', and optionally
|
||||
'cap_drone', 'cap_sat'. Each [B, D] L2-normalized.
|
||||
epoch: Current epoch (0-indexed).
|
||||
total_epochs: Total epochs for temperature schedule.
|
||||
|
||||
Returns:
|
||||
Dict with scalar tensors: 'total', 'img_img', 'sat_cap',
|
||||
'drone_cap', 'cap_cap', plus 'temperature' and 'lambdas'.
|
||||
"""
|
||||
tau = cosine_temperature(
|
||||
epoch=epoch,
|
||||
total_epochs=total_epochs,
|
||||
tau_init=self.temperature_init,
|
||||
tau_final=self.temperature_final,
|
||||
)
|
||||
lambdas = curriculum_lambdas(
|
||||
epoch=epoch,
|
||||
warmup_epochs=self.warmup_epochs,
|
||||
text_ramp_epochs=self.text_ramp_epochs,
|
||||
lambda_ii=self.lambda_ii,
|
||||
lambda_sc_max=self.lambda_sc_max,
|
||||
lambda_dc_max=self.lambda_dc_max,
|
||||
lambda_cc_max=self.lambda_cc_max,
|
||||
)
|
||||
|
||||
drone = embeddings["drone"]
|
||||
sat = embeddings["sat"]
|
||||
|
||||
# Image-image symmetric InfoNCE with asymmetric weights.
|
||||
loss_ii = _symmetric_info_nce(
|
||||
emb_a=drone,
|
||||
emb_b=sat,
|
||||
temperature=tau,
|
||||
label_smoothing=self.label_smoothing,
|
||||
weight_a2b=self.asym_drone_to_sat,
|
||||
weight_b2a=self.asym_sat_to_drone,
|
||||
)
|
||||
|
||||
loss_sc = torch.zeros_like(loss_ii)
|
||||
loss_dc = torch.zeros_like(loss_ii)
|
||||
loss_cc = torch.zeros_like(loss_ii)
|
||||
|
||||
if "cap_sat" in embeddings and lambdas["sat_cap"] > 0:
|
||||
loss_sc = _symmetric_info_nce(
|
||||
emb_a=sat,
|
||||
emb_b=embeddings["cap_sat"],
|
||||
temperature=tau,
|
||||
label_smoothing=self.label_smoothing,
|
||||
)
|
||||
if "cap_drone" in embeddings and lambdas["drone_cap"] > 0:
|
||||
loss_dc = _symmetric_info_nce(
|
||||
emb_a=drone,
|
||||
emb_b=embeddings["cap_drone"],
|
||||
temperature=tau,
|
||||
label_smoothing=self.label_smoothing,
|
||||
)
|
||||
if (
|
||||
"cap_drone" in embeddings
|
||||
and "cap_sat" in embeddings
|
||||
and lambdas["cap_cap"] > 0
|
||||
):
|
||||
loss_cc = _symmetric_info_nce(
|
||||
emb_a=embeddings["cap_drone"],
|
||||
emb_b=embeddings["cap_sat"],
|
||||
temperature=tau,
|
||||
label_smoothing=self.label_smoothing,
|
||||
)
|
||||
|
||||
total = (
|
||||
lambdas["img_img"] * loss_ii
|
||||
+ lambdas["sat_cap"] * loss_sc
|
||||
+ lambdas["drone_cap"] * loss_dc
|
||||
+ lambdas["cap_cap"] * loss_cc
|
||||
)
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"img_img": loss_ii.detach(),
|
||||
"sat_cap": loss_sc.detach(),
|
||||
"drone_cap": loss_dc.detach(),
|
||||
"cap_cap": loss_cc.detach(),
|
||||
"temperature": torch.tensor(tau, device=total.device),
|
||||
"lambda_ii": torch.tensor(lambdas["img_img"], device=total.device),
|
||||
"lambda_sc": torch.tensor(lambdas["sat_cap"], device=total.device),
|
||||
"lambda_dc": torch.tensor(lambdas["drone_cap"], device=total.device),
|
||||
"lambda_cc": torch.tensor(lambdas["cap_cap"], device=total.device),
|
||||
}
|
||||
4
src/models/__init__.py
Normal file
4
src/models/__init__.py
Normal file
@@ -0,0 +1,4 @@
|
||||
"""Model components for caption quality test."""
|
||||
from src.models.dual_encoder import DualEncoderCaptionTest, ProjectionHead
|
||||
|
||||
__all__ = ["DualEncoderCaptionTest", "ProjectionHead"]
|
||||
243
src/models/dual_encoder.py
Normal file
243
src/models/dual_encoder.py
Normal file
@@ -0,0 +1,243 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""Dual encoder for caption quality test on UAV-VisLoc.
|
||||
|
||||
GeoRSCLIP ViT-B/32 backbone (image + text towers, shared 512-dim space).
|
||||
Image encoder is frozen, text encoder has partial unfreeze (last block + projection).
|
||||
Separate trainable projection heads for drone/sat/text branches.
|
||||
"""
|
||||
|
||||
from typing import Literal
|
||||
|
||||
import gin
|
||||
import open_clip
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class ProjectionHead(nn.Module):
|
||||
"""Single-layer L2-normalized projection head.
|
||||
|
||||
Args:
|
||||
in_dim: Input embedding dimension.
|
||||
out_dim: Output embedding dimension (512 for GeoRSCLIP space).
|
||||
use_mlp: If True, use 2-layer MLP with GELU, else Linear.
|
||||
hidden_dim: Hidden dim when use_mlp=True (defaults to 2*in_dim).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_dim: int = 512,
|
||||
out_dim: int = 512,
|
||||
use_mlp: bool = False,
|
||||
hidden_dim: int | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if use_mlp:
|
||||
hidden_dim = hidden_dim or (2 * in_dim)
|
||||
self.proj = nn.Sequential(
|
||||
nn.Linear(in_dim, hidden_dim),
|
||||
nn.GELU(),
|
||||
nn.Linear(hidden_dim, out_dim),
|
||||
)
|
||||
else:
|
||||
self.proj = nn.Linear(in_dim, out_dim)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Project features and L2-normalize.
|
||||
|
||||
Args:
|
||||
x: Input features [B, in_dim].
|
||||
|
||||
Returns:
|
||||
Normalized embeddings [B, out_dim].
|
||||
"""
|
||||
x = self.proj(x)
|
||||
return F.normalize(x, dim=-1)
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class DualEncoderCaptionTest(nn.Module):
|
||||
"""GeoRSCLIP dual encoder for caption quality validation on UAV-VisLoc.
|
||||
|
||||
Shared image encoder for drone and satellite views. Text encoder with
|
||||
partial unfreeze. Three separate trainable projection heads map raw
|
||||
GeoRSCLIP embeddings into the shared 512-dim retrieval space.
|
||||
|
||||
Args:
|
||||
variant: open_clip model variant name (e.g., 'ViT-B-32').
|
||||
pretrained_path: Path to GeoRSCLIP checkpoint (RS5M_ViT-B-32.pt).
|
||||
unfreeze_mode: Which text encoder layers to unfreeze.
|
||||
embed_dim: Output retrieval dimension (default 512).
|
||||
use_mlp_heads: If True, projection heads are 2-layer MLPs.
|
||||
shared_image_head: If True, drone and sat use single projection head.
|
||||
device: torch device.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
variant: str = "ViT-B-32",
|
||||
pretrained_path: str = "RS5M_ViT-B-32.pt",
|
||||
unfreeze_mode: Literal["none", "projection", "last_block", "full"] = "last_block",
|
||||
embed_dim: int = 512,
|
||||
use_mlp_heads: bool = False,
|
||||
shared_image_head: bool = True,
|
||||
device: str = "cuda",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.variant = variant
|
||||
self.embed_dim = embed_dim
|
||||
self.shared_image_head = shared_image_head
|
||||
self.device = device
|
||||
|
||||
# Load open_clip model (GeoRSCLIP compatible with open_clip API).
|
||||
self.model, _, self.preprocess = open_clip.create_model_and_transforms(
|
||||
model_name=variant,
|
||||
pretrained=pretrained_path,
|
||||
device=device,
|
||||
)
|
||||
self.tokenizer = open_clip.get_tokenizer(variant)
|
||||
|
||||
# Native GeoRSCLIP embedding dim (for ViT-B/32 = 512).
|
||||
self._native_dim = self._infer_native_dim()
|
||||
|
||||
# Freeze everything by default.
|
||||
for p in self.model.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
# Apply unfreeze strategy.
|
||||
self._apply_unfreeze(unfreeze_mode)
|
||||
|
||||
# Projection heads (trainable).
|
||||
self.proj_text = ProjectionHead(
|
||||
in_dim=self._native_dim,
|
||||
out_dim=embed_dim,
|
||||
use_mlp=use_mlp_heads,
|
||||
)
|
||||
if shared_image_head:
|
||||
self.proj_image = ProjectionHead(
|
||||
in_dim=self._native_dim,
|
||||
out_dim=embed_dim,
|
||||
use_mlp=use_mlp_heads,
|
||||
)
|
||||
self.proj_drone = None # type: ignore[assignment]
|
||||
self.proj_sat = None # type: ignore[assignment]
|
||||
else:
|
||||
self.proj_image = None # type: ignore[assignment]
|
||||
self.proj_drone = ProjectionHead(
|
||||
in_dim=self._native_dim,
|
||||
out_dim=embed_dim,
|
||||
use_mlp=use_mlp_heads,
|
||||
)
|
||||
self.proj_sat = ProjectionHead(
|
||||
in_dim=self._native_dim,
|
||||
out_dim=embed_dim,
|
||||
use_mlp=use_mlp_heads,
|
||||
)
|
||||
|
||||
def _infer_native_dim(self) -> int:
|
||||
"""Infer native embedding dimension from model (typically 512 for ViT-B/32)."""
|
||||
if hasattr(self.model, "text_projection"):
|
||||
shape = self.model.text_projection.shape
|
||||
return int(shape[1] if shape.ndim == 2 else shape[0])
|
||||
return 512
|
||||
|
||||
def _apply_unfreeze(
|
||||
self,
|
||||
unfreeze_mode: Literal["none", "projection", "last_block", "full"],
|
||||
) -> None:
|
||||
"""Selectively enable gradients for text encoder."""
|
||||
if unfreeze_mode == "none":
|
||||
return
|
||||
if unfreeze_mode == "full":
|
||||
for p in self.model.parameters():
|
||||
p.requires_grad = True
|
||||
return
|
||||
|
||||
# Always unfreeze text_projection if available.
|
||||
if hasattr(self.model, "text_projection"):
|
||||
tp = self.model.text_projection
|
||||
if isinstance(tp, nn.Parameter):
|
||||
tp.requires_grad = True
|
||||
elif isinstance(tp, nn.Module):
|
||||
for p in tp.parameters():
|
||||
p.requires_grad = True
|
||||
|
||||
# Additionally unfreeze last transformer block.
|
||||
if unfreeze_mode == "last_block" and hasattr(self.model, "transformer"):
|
||||
last_block = self.model.transformer.resblocks[-1]
|
||||
for p in last_block.parameters():
|
||||
p.requires_grad = True
|
||||
|
||||
def encode_image(self, images: torch.Tensor) -> torch.Tensor:
|
||||
"""Encode images through GeoRSCLIP image encoder (no projection head).
|
||||
|
||||
Args:
|
||||
images: Preprocessed image tensor [B, 3, H, W].
|
||||
|
||||
Returns:
|
||||
Raw image embeddings [B, native_dim].
|
||||
"""
|
||||
feats = self.model.encode_image(images)
|
||||
return F.normalize(feats, dim=-1)
|
||||
|
||||
def encode_text(self, texts: list[str] | torch.Tensor) -> torch.Tensor:
|
||||
"""Encode text captions through GeoRSCLIP text encoder.
|
||||
|
||||
Args:
|
||||
texts: List of strings or pre-tokenized LongTensor [B, seq_len].
|
||||
|
||||
Returns:
|
||||
Raw text embeddings [B, native_dim].
|
||||
"""
|
||||
if isinstance(texts, (list, tuple)):
|
||||
tokens = self.tokenizer(list(texts)).to(self.device).long()
|
||||
else:
|
||||
tokens = texts.to(self.device).long()
|
||||
feats = self.model.encode_text(tokens)
|
||||
return F.normalize(feats, dim=-1)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
drone_img: torch.Tensor,
|
||||
sat_img: torch.Tensor,
|
||||
caption_drone: list[str] | None = None,
|
||||
caption_sat: list[str] | None = None,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Forward pass producing projected embeddings for all branches.
|
||||
|
||||
Args:
|
||||
drone_img: Drone RGB tensor [B, 3, H, W].
|
||||
sat_img: Satellite RGB tensor [B, 3, H, W].
|
||||
caption_drone: List of drone captions, one per batch item.
|
||||
caption_sat: List of satellite captions, one per batch item.
|
||||
|
||||
Returns:
|
||||
Dict with keys 'drone', 'sat', 'cap_drone', 'cap_sat', each
|
||||
containing [B, embed_dim] L2-normalized embeddings.
|
||||
Keys for missing captions are absent.
|
||||
"""
|
||||
out: dict[str, torch.Tensor] = {}
|
||||
|
||||
drone_feat = self.encode_image(drone_img)
|
||||
sat_feat = self.encode_image(sat_img)
|
||||
|
||||
if self.shared_image_head:
|
||||
out["drone"] = self.proj_image(drone_feat)
|
||||
out["sat"] = self.proj_image(sat_feat)
|
||||
else:
|
||||
out["drone"] = self.proj_drone(drone_feat)
|
||||
out["sat"] = self.proj_sat(sat_feat)
|
||||
|
||||
if caption_drone is not None:
|
||||
out["cap_drone"] = self.proj_text(self.encode_text(caption_drone))
|
||||
if caption_sat is not None:
|
||||
out["cap_sat"] = self.proj_text(self.encode_text(caption_sat))
|
||||
|
||||
return out
|
||||
|
||||
def trainable_parameters(self) -> list[nn.Parameter]:
|
||||
"""Return list of trainable parameters for optimizer construction."""
|
||||
return [p for p in self.parameters() if p.requires_grad]
|
||||
1
src/training/__init__.py
Normal file
1
src/training/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Training loop for caption quality test."""
|
||||
313
src/training/train.py
Normal file
313
src/training/train.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user