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

245
src/datasets/visloc_with_captions.py
View File

@@ -1,26 +1,20 @@
from __future__ import annotations
"""UAV-VisLoc dataset loader augmented with generated captions.
"""UAV-GeoLoc dataset loader with template captions for CVGL caption test.
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]
},
...
]
Reads UAV-GeoLoc Index format (train_query.txt / train_db.txt) and generates
template captions from path metadata (terrain type, scene, altitude, heading).
Captions are produced offline by scripts/generate_captions.py using one of
three strategies: template, VLM, or hybrid (see АНАЛИЗ_caption_quality_test).
train_query.txt format:
Terrain/Volcano/KilaueaVolcano/query/height100_rot315/footage/file.jpeg 0 .../DB/img/crop_X_Y.png ...
Template caption (P3-style fingerprint for cross-view matching):
"Aerial view at 100m facing northwest over volcanic terrain near KilaueaVolcano.
Plan-view features: dark lava flows, crater edges, sparse vegetation patches."
"""
import json
import random
import re
from pathlib import Path
from typing import Any, Callable
@@ -29,130 +23,191 @@ import torch
from PIL import Image
from torch.utils.data import Dataset
# Compass lookup: heading_deg -> direction name.
_COMPASS = ["north", "northeast", "east", "southeast",
"south", "southwest", "west", "northwest"]
# Simple terrain-to-features mapping for template captions.
_TERRAIN_FEATURES: dict[str, str] = {
"Volcano": "lava flows, crater edges, volcanic rock",
"Mountain": "ridgelines, steep slopes, rocky terrain",
"Hill": "rolling hills, gentle slopes, scattered trees",
"Desert": "sand dunes, arid ground, sparse scrub",
"Plain": "flat open fields, agricultural plots, dirt roads",
"Plateau": "flat elevated terrain, cliffs, mesa edges",
"Basin": "lowland depression, dry lake bed, sediment patterns",
"Delta": "river channels, sediment fans, wetland patches",
"Gorge": "deep canyon, exposed rock layers, narrow valley",
"Island": "shoreline, coastal vegetation, water boundary",
"Wetland": "marshes, water channels, aquatic vegetation",
"Glacier": "ice formations, crevasses, glacial moraine",
"Forest": "dense canopy, tree shadows, clearings",
"Farm": "crop fields, irrigation lines, farm buildings",
"Prairie": "grassland, open meadow, fence lines",
"Finca": "rural estate, orchards, scattered structures",
"Calcification": "mineral deposits, white crusts, barren soil",
"StoneForest": "karst pillars, eroded limestone, sparse vegetation",
"Oasis": "palm clusters, water pool, surrounding desert",
"Flowers": "colorful ground cover, floral patterns, open field",
"Terrace": "terraced hillside, stepped fields, retaining walls",
"Snow": "snow cover, tracks, exposed rock patches",
"Pasture": "grazing land, fenced paddocks, grass",
"Danxia": "red sandstone, layered cliffs, erosion patterns",
"Hylare": "sparse woodland, rocky outcrops",
"Karst": "sinkholes, limestone towers, caves",
"Fall": "autumn foliage, colored canopy, leaf litter",
}
# Country/city features.
_COUNTRY_FEATURES: str = "buildings, roads, urban blocks, rooftops, intersections"
def _parse_query_path(query_path: str) -> dict[str, str]:
"""Extract metadata from UAV-GeoLoc query path.
Example: Terrain/Volcano/KilaueaVolcano/query/height100_rot315/footage/file.jpeg
Returns: {category, terrain_type, scene, height_m, heading_deg, heading_dir}
"""
parts = query_path.split("/")
category = parts[0] if parts else "Unknown"
if category == "Terrain":
terrain_type = parts[1] if len(parts) > 1 else "Unknown"
scene = parts[2] if len(parts) > 2 else "Unknown"
elif category == "Country":
terrain_type = "Urban"
scene = "/".join(parts[1:3]) if len(parts) > 2 else "Unknown"
else:
terrain_type = "Unknown"
scene = parts[1] if len(parts) > 1 else "Unknown"
# Parse height and rotation from trajectory folder name.
height_m = "100"
heading_deg = "0"
for part in parts:
m = re.match(r"height(\d+)_rot(\d+)", part)
if m:
height_m = m.group(1)
heading_deg = m.group(2)
break
heading_idx = round(int(heading_deg) / 45) % 8
heading_dir = _COMPASS[heading_idx]
return {
"category": category,
"terrain_type": terrain_type,
"scene": scene,
"height_m": height_m,
"heading_deg": heading_deg,
"heading_dir": heading_dir,
}
def _make_template_caption(meta: dict[str, str]) -> str:
"""Generate a template caption from parsed metadata."""
terrain = meta["terrain_type"]
features = _TERRAIN_FEATURES.get(terrain, _COUNTRY_FEATURES)
return (
f"Aerial view at {meta['height_m']}m facing {meta['heading_dir']} "
f"over {terrain.lower()} terrain near {meta['scene']}. "
f"Plan-view features: {features}."
)
@gin.configurable
class VisLocCaptionDataset(Dataset):
"""UAV-VisLoc pairs with generated captions.
class GeoLocCaptionDataset(Dataset):
"""UAV-GeoLoc pairs with template captions.
Reads train_query.txt, randomly samples one positive crop per query.
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.
query_file: Path to train_query.txt (or test_query.txt).
data_root: Root directory of UAV-GeoLoc dataset.
image_transform: Callable applied to PIL images.
drop_caption_prob: Probability of dropping caption (for ablation).
seed: Random seed.
"""
def __init__(
self,
manifest_path: str,
image_root: str,
query_file: str,
data_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.data_root = Path(data_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.entries: list[dict[str, Any]] = []
self._load_query_file(Path(query_file))
self._validate_entries()
def _load_query_file(self, query_file: Path) -> None:
"""Parse train_query.txt into list of entries."""
with open(query_file) as f:
for line in f:
line = line.strip()
if not line:
continue
parts = line.split()
query_path = parts[0]
# parts[1] is label (always 0), parts[2:] are positive crop paths.
positive_crops = parts[2:]
if not positive_crops:
continue
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}
meta = _parse_query_path(query_path)
caption = _make_template_caption(meta)
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}"
self.entries.append({
"query_path": query_path,
"positive_crops": positive_crops,
"caption": caption,
"meta": meta,
})
def _load_image(self, relative_path: str) -> torch.Tensor:
"""Load image and apply preprocessing."""
path = self.image_root / relative_path
path = self.data_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"])
drone_img = self._load_image(entry["query_path"])
caption_drone = self._maybe_drop(entry[self._caption_key("drone")])
caption_sat = self._maybe_drop(entry[self._caption_key("sat")])
# Randomly sample one positive crop.
crop_path = self._rng.choice(entry["positive_crops"])
sat_img = self._load_image(crop_path)
caption = entry["caption"]
if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
caption = ""
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}"),
"caption_drone": caption,
"pair_id": entry["query_path"],
}
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.
"""
"""Collate into batched dict. Captions stay as string 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],
}

161
src/eval/evaluate.py
View File

@@ -1,9 +1,9 @@
from __future__ import annotations
"""Evaluation utilities for caption quality test.
"""Evaluation 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.
Recall@K for query(drone+text) -> gallery(satellite).
delta_r_at_1 compares caption-aware vs baseline runs.
"""
import json
@@ -23,19 +23,10 @@ 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].
"""
"""Recall@K assuming positives on the diagonal."""
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]
@@ -49,51 +40,29 @@ 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].
"""
"""Encode all samples into query and gallery embeddings."""
model.eval()
all_drone: list[torch.Tensor] = []
all_sat: list[torch.Tensor] = []
all_cap_drone: list[torch.Tensor] = []
all_cap_sat: list[torch.Tensor] = []
all_query: list[torch.Tensor] = []
all_gallery: 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
caption_drone = batch["caption_drone"]
embeddings = model(
drone_img=drone_img,
sat_img=sat_img,
caption_drone=captions_drone,
caption_sat=captions_sat,
caption_drone=caption_drone,
)
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())
all_query.append(embeddings["query"].cpu())
all_gallery.append(embeddings["gallery"].cpu())
out = {
"drone": torch.cat(all_drone, dim=0),
"sat": torch.cat(all_sat, dim=0),
return {
"query": torch.cat(all_query, dim=0),
"gallery": torch.cat(all_gallery, 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(
@@ -101,51 +70,25 @@ def evaluate_retrieval(
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.
"""Compute R@K for query->gallery and gallery->query.
Returns:
Flat dict with keys like 'r@1_drone_to_sat', 'r@5_text_to_sat', etc.
Flat dict: r@1_query_to_gallery, r@5_query_to_gallery, etc.
"""
feats = _encode_dataset(
model=model,
loader=loader,
device=device,
include_captions=include_captions,
)
feats = _encode_dataset(model=model, loader=loader, device=device)
metrics: dict[str, float] = {}
sim_d2s = feats["drone"] @ feats["sat"].t()
sim_s2d = sim_d2s.t()
sim_q2g = feats["query"] @ feats["gallery"].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
for k, val in _recall_at_k(sim_q2g, k_values).items():
metrics[f"r@{k}_query_to_gallery"] = val
for k, val in _recall_at_k(sim_q2g.t(), k_values).items():
metrics[f"r@{k}_gallery_to_query"] = 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
# Gate value for diagnostics.
metrics["gate"] = model.fusion.gate_value
return metrics
@@ -153,64 +96,36 @@ def evaluate_retrieval(
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)}"
)
"""R@1 gain from adding captions: full - baseline."""
key = "r@1_query_to_gallery"
return full_metrics[key] - baseline_metrics[key]
@gin.configurable
def run_evaluation_from_checkpoint(
checkpoint_path: str,
test_manifest: str,
image_root: str,
test_query_file: str,
data_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.
"""
"""Standalone evaluation from checkpoint."""
from src.datasets.visloc_with_captions import (
VisLocCaptionDataset,
GeoLocCaptionDataset,
collate_caption_batch,
)
model = DualEncoderCaptionTest().to(device)
ckpt = torch.load(checkpoint_path, map_location=device)
ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
model.load_state_dict(ckpt["model_state"])
model.eval()
test_ds = VisLocCaptionDataset(
manifest_path=test_manifest,
image_root=image_root,
test_ds = GeoLocCaptionDataset(
query_file=test_query_file,
data_root=data_root,
image_transform=model.preprocess,
)
test_loader = DataLoader(
@@ -222,15 +137,11 @@ def run_evaluation_from_checkpoint(
pin_memory=True,
)
metrics = evaluate_retrieval(
model=model,
loader=test_loader,
device=device,
)
metrics = evaluate_retrieval(model=model, loader=test_loader, device=device)
report = {
"checkpoint": checkpoint_path,
"test_manifest": test_manifest,
"test_query_file": test_query_file,
"metrics": metrics,
}
out = Path(output_path)

208
src/losses/multi_infonce.py
View File

@@ -1,14 +1,10 @@
from __future__ import annotations
"""Multi-term InfoNCE loss for caption quality validation.
"""InfoNCE loss for cross-view geo-localization with optional text fusion.
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).
Single symmetric InfoNCE between query (drone+text fused) and gallery (satellite).
Asymmetric weighting: query->gallery weighted higher (real use-case direction).
Cosine temperature schedule for sharper distribution over training.
"""
import math
@@ -27,26 +23,12 @@ def _symmetric_info_nce(
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.
"""
"""Weighted symmetric InfoNCE. Positives on the diagonal."""
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
@@ -56,85 +38,23 @@ def cosine_temperature(
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.
"""
"""Cosine-decay schedule for InfoNCE temperature."""
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.
class InfoNCELoss(nn.Module):
"""Symmetric InfoNCE with cosine temperature schedule.
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.
temperature_init: Temperature at epoch 0.
temperature_final: Temperature after cosine decay.
label_smoothing: Cross-entropy label smoothing.
weight_q2g: Weight for query->gallery direction.
weight_g2q: Weight for gallery->query direction.
"""
def __init__(
@@ -142,27 +62,15 @@ class MultiTermInfoNCE(nn.Module):
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,
weight_q2g: float = 0.6,
weight_g2q: float = 0.4,
) -> 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
self.weight_q2g = weight_q2g
self.weight_g2q = weight_g2q
def forward(
self,
@@ -170,17 +78,16 @@ class MultiTermInfoNCE(nn.Module):
epoch: int,
total_epochs: int,
) -> dict[str, torch.Tensor]:
"""Compute multi-term loss.
"""Compute InfoNCE loss.
Args:
embeddings: Dict with keys 'drone', 'sat', and optionally
'cap_drone', 'cap_sat'. Each [B, D] L2-normalized.
embeddings: Dict with 'query' and 'gallery' [B, D] L2-normalized,
plus 'gate' (float) from fusion module.
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'.
Dict with 'total', 'temperature', 'gate'.
"""
tau = cosine_temperature(
epoch=epoch,
@@ -188,75 +95,20 @@ class MultiTermInfoNCE(nn.Module):
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,
loss = _symmetric_info_nce(
emb_a=embeddings["query"],
emb_b=embeddings["gallery"],
temperature=tau,
label_smoothing=self.label_smoothing,
weight_a2b=self.asym_drone_to_sat,
weight_b2a=self.asym_sat_to_drone,
weight_a2b=self.weight_q2g,
weight_b2a=self.weight_g2q,
)
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
)
gate = embeddings.get("gate", 1.0)
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),
"total": loss,
"temperature": torch.tensor(tau, device=loss.device),
"gate": torch.tensor(gate, device=loss.device),
}

222
src/models/dual_encoder.py
View File

@@ -1,10 +1,16 @@
from __future__ import annotations
"""Dual encoder for caption quality test on UAV-VisLoc.
"""Dual encoder for caption quality test on cross-view geo-localization.
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.
Image encoder frozen. Text encoder with partial unfreeze.
Architecture:
Query branch: GeoRSCLIP_img(drone) + GeoRSCLIP_text(caption) -> GatedFusion -> proj -> query_emb
Gallery branch: GeoRSCLIP_img(sat) -> proj -> gallery_emb
Loss: InfoNCE(query_emb, gallery_emb)
Baseline mode: fusion gate forced to 1.0 (text ignored).
"""
from typing import Literal
@@ -16,16 +22,8 @@ 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).
"""
"""MLP projection head with L2 normalization."""
def __init__(
self,
@@ -46,33 +44,56 @@ class ProjectionHead(nn.Module):
self.proj = nn.Linear(in_dim, out_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Project features and L2-normalize.
return F.normalize(self.proj(x), dim=-1)
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 GatedFusion(nn.Module):
"""Learnable gated fusion of image and text embeddings.
q = sigma(alpha) * img + (1 - sigma(alpha)) * text
alpha is a single learnable scalar, initialized so that gate ~ init_gate.
When baseline_mode=True, gate is clamped to 1.0 (text contribution = 0).
"""
def __init__(self, init_gate: float = 0.7, baseline_mode: bool = False) -> None:
super().__init__()
# alpha is in logit space: sigmoid(alpha) = init_gate
init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
self.alpha = nn.Parameter(init_alpha)
self.baseline_mode = baseline_mode
def forward(
self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
if text_feat is None or self.baseline_mode:
return img_feat
gate = torch.sigmoid(self.alpha)
return gate * img_feat + (1.0 - gate) * text_feat
@property
def gate_value(self) -> float:
"""Current gate value (image weight). 1.0 = text ignored."""
if self.baseline_mode:
return 1.0
return torch.sigmoid(self.alpha).item()
@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.
"""GeoRSCLIP dual encoder with gated text fusion on query branch.
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.
variant: open_clip model variant name.
pretrained_path: Path to GeoRSCLIP checkpoint.
unfreeze_mode: Text encoder unfreeze strategy.
embed_dim: Output retrieval embedding dimension.
use_mlp_heads: Use 2-layer MLP projection heads.
baseline_mode: If True, fusion gate = 1.0 (no text).
init_gate: Initial gate value (image weight).
device: torch device.
"""
@@ -80,19 +101,19 @@ class DualEncoderCaptionTest(nn.Module):
self,
variant: str = "ViT-B-32",
pretrained_path: str = "RS5M_ViT-B-32.pt",
unfreeze_mode: Literal["none", "projection", "last_block", "full"] = "last_block",
unfreeze_mode: Literal["none", "projection", "last_block"] = "last_block",
embed_dim: int = 512,
use_mlp_heads: bool = False,
shared_image_head: bool = True,
baseline_mode: bool = False,
init_gate: float = 0.7,
device: str = "cuda",
) -> None:
super().__init__()
self.variant = variant
self.embed_dim = embed_dim
self.shared_image_head = shared_image_head
self.device = device
self.baseline_mode = baseline_mode
# Load open_clip model (GeoRSCLIP compatible with open_clip API).
# Load GeoRSCLIP via open_clip.
self.model, _, self.preprocess = open_clip.create_model_and_transforms(
model_name=variant,
pretrained=pretrained_path,
@@ -100,45 +121,28 @@ class DualEncoderCaptionTest(nn.Module):
)
self.tokenizer = open_clip.get_tokenizer(variant)
# Native GeoRSCLIP embedding dim (for ViT-B/32 = 512).
self._native_dim = self._infer_native_dim()
native_dim = self._infer_native_dim()
# Freeze everything by default.
# Freeze everything.
for p in self.model.parameters():
p.requires_grad = False
# Apply unfreeze strategy.
self._apply_unfreeze(unfreeze_mode)
# Selectively unfreeze text encoder (only if not baseline).
if not baseline_mode:
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,
# Gated fusion on query branch.
self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
# Projection heads.
self.proj_query = ProjectionHead(
in_dim=native_dim, out_dim=embed_dim, use_mlp=use_mlp_heads,
)
self.proj_gallery = ProjectionHead(
in_dim=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])
@@ -146,17 +150,11 @@ class DualEncoderCaptionTest(nn.Module):
def _apply_unfreeze(
self,
unfreeze_mode: Literal["none", "projection", "last_block", "full"],
unfreeze_mode: Literal["none", "projection", "last_block"],
) -> 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.
# Unfreeze text_projection.
if hasattr(self.model, "text_projection"):
tp = self.model.text_projection
if isinstance(tp, nn.Parameter):
@@ -164,38 +162,17 @@ class DualEncoderCaptionTest(nn.Module):
elif isinstance(tp, nn.Module):
for p in tp.parameters():
p.requires_grad = True
# Additionally unfreeze last transformer block.
# 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():
for p in self.model.transformer.resblocks[-1].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()
def encode_text(self, texts: list[str]) -> torch.Tensor:
tokens = self.tokenizer(list(texts)).to(self.device).long()
feats = self.model.encode_text(tokens)
return F.normalize(feats, dim=-1)
@@ -204,40 +181,37 @@ class DualEncoderCaptionTest(nn.Module):
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.
"""Forward pass.
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.
drone_img: Drone images [B, 3, H, W].
sat_img: Satellite images [B, 3, H, W].
caption_drone: Drone captions (P3 fingerprint), one per sample.
Returns:
Dict with keys 'drone', 'sat', 'cap_drone', 'cap_sat', each
containing [B, embed_dim] L2-normalized embeddings.
Keys for missing captions are absent.
Dict with 'query' [B, embed_dim], 'gallery' [B, embed_dim],
and 'gate' (scalar) for logging.
"""
out: dict[str, torch.Tensor] = {}
drone_feat = self.encode_image(drone_img)
# Gallery branch: satellite only.
sat_feat = self.encode_image(sat_img)
gallery = self.proj_gallery(sat_feat)
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)
# Query branch: drone + optional text fusion.
drone_feat = self.encode_image(drone_img)
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))
text_feat = None
if caption_drone is not None and not self.baseline_mode:
text_feat = self.encode_text(caption_drone)
return out
fused = self.fusion(drone_feat, text_feat)
query = self.proj_query(fused)
return {
"query": query,
"gallery": gallery,
"gate": self.fusion.gate_value,
}
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]

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)