1061 lines
45 KiB
Python
1061 lines
45 KiB
Python
from __future__ import annotations
|
|
|
|
"""Trainer for CVGL caption test on GTA-UAV-LR.
|
|
|
|
Decomposed from src/training/train_gtauav.py::train into a class with one
|
|
orchestrating method `train()` plus dedicated `_setup_*` / `_build_*` /
|
|
`_train_*` / `_evaluate_*` methods.
|
|
|
|
Lifecycle:
|
|
Trainer(...) → train() → done.
|
|
|
|
`train()` calls _build_* in dependency order, then _train_loop, then
|
|
_final_evaluation; cleanup is in a `finally` block.
|
|
|
|
Currently supports DINOv3 and StripNet backbones only. SOFIA v1/v7.1 model
|
|
classes live in src/models/sofia_v1/ and src/models/sofia_v71/ but are not
|
|
yet wired into this training pipeline (no caption-aware fusion encoder
|
|
wrapper exists for them). Their gin presets remain in in/config_files/
|
|
for future integration; loading one will fail at config_loader level.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import math
|
|
import time
|
|
import warnings
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.amp import GradScaler, autocast
|
|
from torch.optim import AdamW
|
|
from torch.optim.lr_scheduler import LambdaLR
|
|
from torch.utils.data import DataLoader
|
|
from tqdm import tqdm
|
|
|
|
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.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler
|
|
from src.datasets.embedding_cache import EmbeddingCache
|
|
from src.datasets.gtauav_dataset import (
|
|
GTAUAVDataset,
|
|
GTAUAVDroneQuery,
|
|
collate_drone_query,
|
|
collate_gtauav_batch,
|
|
)
|
|
from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler
|
|
from src.eval.evaluator import evaluate
|
|
from src.losses.hard_negatives import NegativeMemoryBank
|
|
from src.losses.multi_infonce import InfoNCELoss
|
|
from src.losses.weighted_infonce import WeightedInfoNCELoss
|
|
from src.models.asymmetric_encoder import (
|
|
AsymmetricEncoder,
|
|
get_dino_transform,
|
|
get_drone_train_transform,
|
|
get_satellite_train_transform,
|
|
)
|
|
from src.training.csv_logger import CSVLogger
|
|
from src.training.grad_monitor import compute_gradient_norms, log_gradient_summary
|
|
from src.training.plot_metrics import generate_plots
|
|
from src.training.profiling import TrainingProfiler, print_model_summary
|
|
from src.training.trackers import ExperimentTracker
|
|
from src.utils.io_utils import atomic_save_torch, clear_vram
|
|
from src.utils.seed_utils import set_seed
|
|
|
|
LOGGER = logging.getLogger("caption_test.trainer")
|
|
|
|
# Type alias for the family-specific models config.
|
|
# SOFIA v1/v71 will join this union once their fusion encoders are written.
|
|
ModelsConfig = DINOv3ModelsConfig | StripNetModelsConfig
|
|
|
|
# Backbones currently wired into this trainer.
|
|
_SUPPORTED_BACKBONES: frozenset[str] = frozenset({"dinov3", "stripnet"})
|
|
|
|
|
|
def _build_param_groups(
|
|
model: AsymmetricEncoder,
|
|
lr: float,
|
|
text_lr_factor: float,
|
|
stripnet_backbone_lr_factor: float = 0.1,
|
|
) -> list[dict]:
|
|
"""Build parameter groups with separate LR for text encoder and StripNet backbone.
|
|
|
|
Group 0: projections + heads + MONA + (logit_scale appended later).
|
|
Group 1: DGTRS-CLIP text encoder (lr * text_lr_factor).
|
|
Group 2 (optional): StripNet backbone when unfrozen (lr * stripnet_backbone_lr_factor).
|
|
"""
|
|
main_params: list[nn.Parameter] = []
|
|
text_params: list[nn.Parameter] = []
|
|
stripnet_backbone_params: list[nn.Parameter] = []
|
|
|
|
for name, p in model.named_parameters():
|
|
if not p.requires_grad:
|
|
continue
|
|
if "text_encoder" in name:
|
|
text_params.append(p)
|
|
elif name.startswith("backbone.") or name.startswith("stripnet."):
|
|
stripnet_backbone_params.append(p)
|
|
else:
|
|
main_params.append(p)
|
|
|
|
groups: list[dict] = [{"params": main_params, "lr": lr, "name": "main"}]
|
|
if text_params:
|
|
groups.append({"params": text_params, "lr": lr * text_lr_factor, "name": "text"})
|
|
if stripnet_backbone_params:
|
|
groups.append({
|
|
"params": stripnet_backbone_params,
|
|
"lr": lr * stripnet_backbone_lr_factor,
|
|
"name": "stripnet_backbone",
|
|
})
|
|
return groups
|
|
|
|
|
|
def _cosine_warmup_schedule(warmup_steps: int, total_steps: int):
|
|
"""Return a lr_lambda for LambdaLR: linear warmup + cosine decay."""
|
|
|
|
def lr_lambda(step: int) -> float:
|
|
if step < warmup_steps:
|
|
return float(step) / max(1, warmup_steps)
|
|
progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
|
|
return 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0)))
|
|
|
|
return lr_lambda
|
|
|
|
@torch.no_grad()
|
|
def _embed_drone_queries(
|
|
model: AsymmetricEncoder,
|
|
train_ds: GTAUAVDataset,
|
|
device: str,
|
|
batch_size: int,
|
|
num_workers: int,
|
|
) -> torch.Tensor:
|
|
"""Forward all drone queries and return [N, D] embeddings on CPU.
|
|
|
|
Used by DynamicSimilaritySampler to rank drones by visual similarity.
|
|
Runs with model.eval() but restores original train state afterwards.
|
|
"""
|
|
was_training = model.training
|
|
model.eval()
|
|
|
|
query_ds = GTAUAVDroneQuery(train_ds)
|
|
loader = DataLoader(
|
|
query_ds,
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
num_workers=num_workers,
|
|
collate_fn=collate_drone_query,
|
|
pin_memory=True,
|
|
)
|
|
|
|
embs: list[torch.Tensor] = []
|
|
for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False):
|
|
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
|
q = model.encode_query(
|
|
drone_img,
|
|
batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
|
|
)
|
|
embs.append(q.cpu())
|
|
|
|
if was_training:
|
|
model.train()
|
|
return torch.cat(embs, dim=0)
|
|
|
|
|
|
class Trainer:
|
|
"""Orchestrates one training run.
|
|
|
|
All gin parameters arrive as 6 config objects; runtime state (model,
|
|
optimizer, loaders, ...) is built lazily by _build_* methods and lives
|
|
on `self`. `train()` calls them in dependency order.
|
|
|
|
Backbones supported: 'dinov3', 'stripnet'.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pipeline_cfg: PipelineConfig,
|
|
hardware_cfg: HardwareConfig,
|
|
training_cfg: TrainingConfig,
|
|
tracking_cfg: TrackingConfig,
|
|
models_common_cfg: ModelsCommonConfig,
|
|
models_cfg: ModelsConfig,
|
|
) -> None:
|
|
self.pipeline_cfg = pipeline_cfg
|
|
self.hardware_cfg = hardware_cfg
|
|
self.training_cfg = training_cfg
|
|
self.tracking_cfg = tracking_cfg
|
|
self.models_common_cfg = models_common_cfg
|
|
self.models_cfg = models_cfg
|
|
|
|
# Runtime state — populated by _build_* methods.
|
|
self.output_dir: Path | None = None
|
|
self.full_config: dict | None = None
|
|
self.tracker: ExperimentTracker | None = None
|
|
self.csv_logger: CSVLogger | None = None
|
|
self.model: nn.Module | None = None
|
|
self.loss_fn: nn.Module | None = None
|
|
self.neg_bank: NegativeMemoryBank | None = None
|
|
self.optimizer: torch.optim.Optimizer | None = None
|
|
self.scheduler: LambdaLR | None = None
|
|
self.scaler: GradScaler | None = None
|
|
self.train_ds: GTAUAVDataset | None = None
|
|
self.test_ds: GTAUAVDataset | None = None
|
|
self.train_eval_ds: GTAUAVDataset | None = None
|
|
self.train_loader: DataLoader | None = None
|
|
self.test_loader: DataLoader | None = None
|
|
self.train_eval_loader: DataLoader | None = None
|
|
self.batch_sampler: DynamicSimilaritySampler | MutuallyExclusiveSampler | None = None
|
|
self.emb_cache: EmbeddingCache | None = None
|
|
self.profiler: TrainingProfiler | None = None
|
|
self.resume_ckpt: dict | None = None
|
|
|
|
# Loop state.
|
|
self.start_epoch: int = 0
|
|
self.global_step: int = 0
|
|
self.best_r1: float = 0.0
|
|
self.history: list[dict] = []
|
|
self.steps_per_epoch: int = 0
|
|
|
|
# ===================================================================
|
|
# Public entry point
|
|
# ===================================================================
|
|
|
|
def train(self) -> None:
|
|
"""Full pipeline: setup → build → train → evaluate → cleanup."""
|
|
self._validate_backbone()
|
|
clear_vram()
|
|
set_seed(self.pipeline_cfg.seed)
|
|
self._setup_output_dir()
|
|
self._setup_tracker()
|
|
self._build_model()
|
|
self._configure_gradient_checkpointing()
|
|
self._log_model_summary()
|
|
self._build_loss()
|
|
self._build_neg_bank()
|
|
self._build_data_loaders()
|
|
self._build_optimizer_and_scheduler()
|
|
self._restore_from_resume()
|
|
self._setup_profiler()
|
|
|
|
#! ------ passed: OK --------------
|
|
try:
|
|
self._train_loop()
|
|
self._final_evaluation()
|
|
finally:
|
|
self._cleanup()
|
|
|
|
# ===================================================================
|
|
# Build phase
|
|
# ===================================================================
|
|
|
|
def _validate_backbone(self) -> None:
|
|
"""Reject unsupported backbones up front with a helpful message."""
|
|
LOGGER.info("⚙️ Validate backbone")
|
|
backbone = self.models_common_cfg.backbone
|
|
if backbone not in _SUPPORTED_BACKBONES:
|
|
raise NotImplementedError(
|
|
f"Trainer does not support backbone={backbone!r} yet. "
|
|
f"Supported backbones: {sorted(_SUPPORTED_BACKBONES)}. "
|
|
f"SOFIA v1/v7.1 model classes exist in src/models/sofia_v1/ and "
|
|
f"src/models/sofia_v71/, but a caption-aware fusion encoder "
|
|
f"wrapper has not been written for them. To enable, create "
|
|
f"the wrapper class with .encode_query/.encode_gallery/"
|
|
f".fusion_query/.fusion_gallery API matching AsymmetricEncoder, "
|
|
f"then add the corresponding branch to _build_model.",
|
|
)
|
|
|
|
def _setup_output_dir(self) -> None:
|
|
"""Create output_dir, save config.json, init csv_logger."""
|
|
LOGGER.info("⚙️ Setup out dir")
|
|
self.output_dir = Path(self.pipeline_cfg.output_dir)
|
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Merge all 6 config objects into one dict for full traceability.
|
|
self.full_config = {
|
|
"pipeline": vars(self.pipeline_cfg),
|
|
"hardware": vars(self.hardware_cfg),
|
|
"training": vars(self.training_cfg),
|
|
"tracking": vars(self.tracking_cfg),
|
|
"models_common": vars(self.models_common_cfg),
|
|
"models": vars(self.models_cfg),
|
|
}
|
|
with (self.output_dir / "config.json").open("w") as f:
|
|
json.dump(self.full_config, f, indent=2)
|
|
|
|
self.csv_logger = CSVLogger(self.output_dir)
|
|
|
|
def _setup_tracker(self) -> None:
|
|
"""W&B + TensorBoard tracker."""
|
|
LOGGER.info("⚙️ Setup tracker...")
|
|
assert self.output_dir is not None and self.full_config is not None
|
|
self.tracker = ExperimentTracker(
|
|
output_dir=self.output_dir,
|
|
config=self.full_config,
|
|
use_wandb=self.tracking_cfg.use_wandb,
|
|
use_tb=self.tracking_cfg.use_tb,
|
|
wandb_project=self.tracking_cfg.wandb_project,
|
|
wandb_run_name=self.tracking_cfg.wandb_run_name,
|
|
wandb_entity=self.tracking_cfg.wandb_entity,
|
|
)
|
|
|
|
def _build_model(self) -> None:
|
|
"""Build (or load) the encoder model based on the active backbone."""
|
|
LOGGER.info("⚙️ Build model...")
|
|
backbone = self.models_common_cfg.backbone
|
|
|
|
if self.pipeline_cfg.resume_from is not None:
|
|
self._build_model_from_resume(backbone)
|
|
return
|
|
|
|
# Fresh build.
|
|
mode_str = "baseline (no text)" if self.models_common_cfg.baseline_mode else "with text (L1/L2/L3)"
|
|
if backbone == "stripnet":
|
|
enc_str = "StripNet-small (shared, 512→1024 proj)"
|
|
else: # dinov3
|
|
assert isinstance(self.models_cfg, DINOv3ModelsConfig)
|
|
enc_str = "shared DINOv3 WEB" if self.models_cfg.shared_encoder \
|
|
else "asymmetric (WEB + SAT)"
|
|
LOGGER.info("Building model — %s, %s", mode_str, enc_str)
|
|
|
|
if backbone == "stripnet":
|
|
self.model = self._build_stripnet_model()
|
|
else: # dinov3
|
|
self.model = self._build_dinov3_model()
|
|
LOGGER.info("embed_dim=%d", self.model.embed_dim)
|
|
|
|
def _build_model_from_resume(self, backbone: str) -> None:
|
|
"""Resume model from checkpoint. Sets self.model, self.resume_ckpt, self.start_epoch."""
|
|
LOGGER.info("⚙️ Build model from resume...")
|
|
LOGGER.info("Resuming from %s", self.pipeline_cfg.resume_from)
|
|
# Both DINOv3 and StripNet go through AsymmetricEncoder.load_checkpoint.
|
|
# Note: load_checkpoint doesn't support StripNet — known existing limitation.
|
|
if isinstance(self.models_cfg, DINOv3ModelsConfig):
|
|
dino_web_path = self.models_cfg.dino_web_path
|
|
dino_sat_path = self.models_cfg.dino_sat_path
|
|
else:
|
|
# StripNet preset on resume — fall back to original defaults.
|
|
dino_web_path = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth"
|
|
dino_sat_path = "nn_models/DINO_SAT/model.safetensors"
|
|
self.model, self.resume_ckpt = AsymmetricEncoder.load_checkpoint(
|
|
self.pipeline_cfg.resume_from,
|
|
dino_web_path=dino_web_path,
|
|
dino_sat_path=dino_sat_path,
|
|
lrsclip_path=self.models_common_cfg.lrsclip_path,
|
|
device=self.hardware_cfg.device,
|
|
)
|
|
self.start_epoch = self.resume_ckpt.get("epoch", -1) + 1
|
|
|
|
def _build_stripnet_model(self) -> nn.Module:
|
|
"""Construct AsymmetricEncoder configured for StripNet."""
|
|
LOGGER.info("⚙️ Build StripNet model...")
|
|
assert isinstance(self.models_cfg, StripNetModelsConfig)
|
|
m = self.models_cfg
|
|
# DINO paths passed but ignored at runtime when backbone='stripnet'.
|
|
# mona_bottleneck=64 matches the original TrainConfigGTAUAV default —
|
|
# used by inject_conv_mona_into_stripnet.
|
|
return AsymmetricEncoder(
|
|
dino_web_path="nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
|
|
dino_sat_path="nn_models/DINO_SAT/model.safetensors",
|
|
lrsclip_path=self.models_common_cfg.lrsclip_path,
|
|
init_gate=self.models_common_cfg.init_gate,
|
|
baseline_mode=self.models_common_cfg.baseline_mode,
|
|
shared_encoder=True, # StripNet always shared (overridden internally)
|
|
mona_bottleneck=64, # matches old TrainConfigGTAUAV.mona_bottleneck
|
|
mona_last_n_blocks=12, # not used for StripNet, but accepted by sig
|
|
device=self.hardware_cfg.device,
|
|
backbone="stripnet",
|
|
stripnet_path=m.stripnet_path,
|
|
stripnet_mona_last_n_stages=m.stripnet_mona_last_n_stages,
|
|
stripnet_freeze=m.stripnet_freeze,
|
|
).to(self.hardware_cfg.device)
|
|
|
|
def _build_dinov3_model(self) -> nn.Module:
|
|
LOGGER.info("⚙️ Build DINOv3 model...")
|
|
"""Construct AsymmetricEncoder configured for DINOv3."""
|
|
assert isinstance(self.models_cfg, DINOv3ModelsConfig)
|
|
m = self.models_cfg
|
|
# stripnet_path passed with the original default — ignored at runtime
|
|
# for DINOv3 backbone.
|
|
return AsymmetricEncoder(
|
|
dino_web_path=m.dino_web_path,
|
|
dino_sat_path=m.dino_sat_path,
|
|
lrsclip_path=self.models_common_cfg.lrsclip_path,
|
|
init_gate=self.models_common_cfg.init_gate,
|
|
baseline_mode=self.models_common_cfg.baseline_mode,
|
|
shared_encoder=m.shared_encoder,
|
|
mona_bottleneck=m.mona_bottleneck,
|
|
mona_last_n_blocks=m.mona_last_n_blocks,
|
|
device=self.hardware_cfg.device,
|
|
backbone="dinov3",
|
|
stripnet_path="nn_models/STRIPNET/stripnet_s.pth",
|
|
stripnet_mona_last_n_stages=0,
|
|
stripnet_freeze=True,
|
|
).to(self.hardware_cfg.device)
|
|
|
|
def _configure_gradient_checkpointing(self) -> None:
|
|
"""Enable gradient checkpointing on encoders that support it."""
|
|
LOGGER.info("⚙️ Configure gradient checkpointing...")
|
|
assert self.model is not None
|
|
backbone = self.models_common_cfg.backbone
|
|
if not self.hardware_cfg.gradient_checkpointing:
|
|
return
|
|
if backbone == "dinov3":
|
|
assert isinstance(self.models_cfg, DINOv3ModelsConfig)
|
|
if self.models_cfg.shared_encoder:
|
|
self.model.image_encoder.set_gradient_checkpointing(True)
|
|
else:
|
|
self.model.drone_encoder.set_gradient_checkpointing(True)
|
|
self.model.sat_encoder.set_gradient_checkpointing(True)
|
|
if self.model.text_encoder is not None:
|
|
self.model.text_encoder.transformer.gradient_checkpointing = True
|
|
LOGGER.info("✅ Gradient checkpointing enabled (DINOv3 + DGTRS)")
|
|
elif backbone == "stripnet":
|
|
if self.model.text_encoder is not None:
|
|
self.model.text_encoder.transformer.gradient_checkpointing = True
|
|
LOGGER.info("✅ Gradient checkpointing enabled (DGTRS only; StripNet doesn't support it)")
|
|
|
|
def _log_model_summary(self) -> None:
|
|
"""Log trainable param count, save model_summary.txt, hook W&B."""
|
|
assert self.model is not None and self.output_dir is not None and self.tracker is not None
|
|
n_trainable = sum(p.numel() for p in self.model.trainable_parameters())
|
|
n_total = sum(p.numel() for p in self.model.parameters())
|
|
LOGGER.info(
|
|
"trainable=%s (%.2f%%) total=%s",
|
|
f"{n_trainable:,}", 100.0 * n_trainable / max(n_total, 1), f"{n_total:,}",
|
|
)
|
|
model_summary = print_model_summary(self.model, device=self.hardware_cfg.device)
|
|
(self.output_dir / "model_summary.txt").write_text(model_summary)
|
|
if self.tracker.has_wandb:
|
|
self.tracker.watch_model(self.model, log_freq=50)
|
|
|
|
def _build_loss(self) -> None:
|
|
"""Build InfoNCELoss or WeightedInfoNCELoss based on training_cfg.loss_type."""
|
|
LOGGER.info("⚙️ Build loss...")
|
|
t = self.training_cfg
|
|
if t.loss_type == "symmetric":
|
|
self.loss_fn = InfoNCELoss(
|
|
temperature_init=t.tau_init,
|
|
temperature_final=t.tau_final,
|
|
label_smoothing=t.label_smoothing,
|
|
weight_q2g=t.weight_q2g,
|
|
weight_g2q=t.weight_g2q,
|
|
learnable_temperature=t.learnable_temperature,
|
|
tau_min=t.tau_min,
|
|
tau_max=t.tau_max,
|
|
hard_mining_k=t.hard_mining_k,
|
|
)
|
|
loss_name = "SymmetricInfoNCE"
|
|
elif t.loss_type == "weighted":
|
|
self.loss_fn = WeightedInfoNCELoss(
|
|
temperature_init=t.tau_init,
|
|
learnable_temperature=t.learnable_temperature,
|
|
label_smoothing=t.label_smoothing,
|
|
k=t.weighted_loss_k,
|
|
tau_min=t.tau_min,
|
|
tau_max=t.tau_max,
|
|
)
|
|
loss_name = "WeightedInfoNCE"
|
|
else:
|
|
raise ValueError(
|
|
f"Unknown loss_type={t.loss_type!r} (expected 'symmetric' or 'weighted')",
|
|
)
|
|
LOGGER.info(
|
|
"Loss: %s Temperature: %s (init=%.3f) q2g=%.2f g2q=%.2f",
|
|
loss_name,
|
|
"learnable" if t.learnable_temperature else "fixed",
|
|
t.tau_init, t.weight_q2g, t.weight_g2q,
|
|
)
|
|
|
|
def _build_neg_bank(self) -> None:
|
|
"""Optional NegativeMemoryBank for hard-negative mining."""
|
|
LOGGER.info("⚙️ Build negative bank...")
|
|
assert self.model is not None
|
|
if self.training_cfg.neg_bank_size > 0:
|
|
self.neg_bank = NegativeMemoryBank(
|
|
size=self.training_cfg.neg_bank_size,
|
|
dim=self.model.embed_dim,
|
|
).to(self.hardware_cfg.device)
|
|
LOGGER.info(
|
|
"Negative memory bank: size=%d, dim=%d",
|
|
self.training_cfg.neg_bank_size, self.model.embed_dim,
|
|
)
|
|
|
|
def _build_data_loaders(self) -> None:
|
|
"""Build train/test/train_eval datasets, samplers, loaders."""
|
|
LOGGER.info("⚙️ Build dataloaders...")
|
|
drone_train_tf = get_drone_train_transform(image_size=256)
|
|
sat_train_tf = get_satellite_train_transform(image_size=256)
|
|
eval_tf = get_dino_transform(image_size=256)
|
|
|
|
self.train_ds = GTAUAVDataset(
|
|
pair_json=self.pipeline_cfg.train_json,
|
|
rgb_root=self.pipeline_cfg.rgb_root,
|
|
caption_root=self.pipeline_cfg.caption_root,
|
|
drone_transform=drone_train_tf,
|
|
sat_transform=sat_train_tf,
|
|
filter_meta=self.pipeline_cfg.filter_meta,
|
|
)
|
|
self.test_ds = GTAUAVDataset(
|
|
pair_json=self.pipeline_cfg.test_json,
|
|
rgb_root=self.pipeline_cfg.rgb_root,
|
|
caption_root=self.pipeline_cfg.caption_root,
|
|
image_transform=eval_tf,
|
|
filter_meta=self.pipeline_cfg.filter_meta,
|
|
)
|
|
self.train_eval_ds = GTAUAVDataset(
|
|
pair_json=self.pipeline_cfg.train_json,
|
|
rgb_root=self.pipeline_cfg.rgb_root,
|
|
caption_root=self.pipeline_cfg.caption_root,
|
|
image_transform=eval_tf,
|
|
filter_meta=self.pipeline_cfg.filter_meta,
|
|
)
|
|
|
|
self._build_batch_sampler()
|
|
|
|
if self.batch_sampler is not None:
|
|
self.train_loader = DataLoader(
|
|
self.train_ds,
|
|
batch_sampler=self.batch_sampler,
|
|
num_workers=self.hardware_cfg.num_workers,
|
|
collate_fn=collate_gtauav_batch,
|
|
pin_memory=True,
|
|
)
|
|
else:
|
|
self.train_loader = DataLoader(
|
|
self.train_ds,
|
|
batch_size=self.hardware_cfg.batch_size,
|
|
shuffle=True,
|
|
num_workers=self.hardware_cfg.num_workers,
|
|
collate_fn=collate_gtauav_batch,
|
|
pin_memory=True,
|
|
drop_last=True,
|
|
)
|
|
|
|
if self.training_cfg.dss_cache_dir is not None:
|
|
self.emb_cache = EmbeddingCache(self.training_cfg.dss_cache_dir)
|
|
LOGGER.info("DSS embedding cache: %s", self.training_cfg.dss_cache_dir)
|
|
|
|
self.test_loader = DataLoader(
|
|
self.test_ds,
|
|
batch_size=self.hardware_cfg.batch_size,
|
|
shuffle=False,
|
|
num_workers=self.hardware_cfg.num_workers,
|
|
collate_fn=collate_gtauav_batch,
|
|
pin_memory=True,
|
|
)
|
|
self.train_eval_loader = DataLoader(
|
|
self.train_eval_ds,
|
|
batch_size=self.hardware_cfg.batch_size,
|
|
shuffle=False,
|
|
num_workers=self.hardware_cfg.num_workers,
|
|
collate_fn=collate_gtauav_batch,
|
|
pin_memory=True,
|
|
)
|
|
|
|
effective_batch = self.hardware_cfg.batch_size * self.hardware_cfg.grad_accum_steps
|
|
LOGGER.info(
|
|
"train=%d test=%d batch=%d accum=%d effective_batch=%d",
|
|
len(self.train_ds), len(self.test_ds),
|
|
self.hardware_cfg.batch_size, self.hardware_cfg.grad_accum_steps, effective_batch,
|
|
)
|
|
|
|
def _build_batch_sampler(self) -> None:
|
|
"""Choose between DSS / mutex / plain shuffle samplers."""
|
|
assert self.train_ds is not None
|
|
sat_cand_list = [entry["sat_candidates"] for entry in self.train_ds.entries]
|
|
# Backward compat alias.
|
|
effective_sampler_type = (
|
|
self.training_cfg.sampler_type
|
|
if self.training_cfg.use_mutex_sampler else "none"
|
|
)
|
|
t = self.training_cfg
|
|
if effective_sampler_type == "dss":
|
|
self.batch_sampler = DynamicSimilaritySampler(
|
|
sat_cand_list,
|
|
batch_size=self.hardware_cfg.batch_size,
|
|
shuffle=True,
|
|
seed=self.pipeline_cfg.seed,
|
|
knn_device=t.dss_knn_device,
|
|
use_lsh=t.dss_use_lsh,
|
|
lsh_num_tables=t.dss_lsh_num_tables,
|
|
lsh_num_bits=t.dss_lsh_num_bits,
|
|
)
|
|
LOGGER.info(
|
|
"Sampler: DynamicSimilarity — kNN on %s%s, warmup=%d, re-embed every %d epochs",
|
|
t.dss_knn_device, " + LSH" if t.dss_use_lsh else "",
|
|
t.dss_warmup_epochs, t.dss_reembed_every,
|
|
)
|
|
elif effective_sampler_type == "mutex":
|
|
self.batch_sampler = MutuallyExclusiveSampler(
|
|
sat_cand_list,
|
|
batch_size=self.hardware_cfg.batch_size,
|
|
shuffle=True,
|
|
seed=self.pipeline_cfg.seed,
|
|
)
|
|
LOGGER.info("Sampler: MutuallyExclusive — no false negatives within a batch")
|
|
else:
|
|
self.batch_sampler = None
|
|
LOGGER.info("Sampler: default shuffle (no mutex / no DSS)")
|
|
|
|
def _build_optimizer_and_scheduler(self) -> None:
|
|
"""Build AdamW with per-group LR + cosine-warmup scheduler + GradScaler."""
|
|
LOGGER.info("⚙️ Build optimizer & scheduler...")
|
|
assert self.model is not None and self.loss_fn is not None and self.train_loader is not None
|
|
t = self.training_cfg
|
|
|
|
stripnet_lr_factor = (
|
|
self.models_cfg.stripnet_backbone_lr_factor
|
|
if isinstance(self.models_cfg, StripNetModelsConfig)
|
|
else 0.1
|
|
)
|
|
param_groups = _build_param_groups(
|
|
self.model,
|
|
t.learning_rate,
|
|
t.text_lr_factor,
|
|
stripnet_backbone_lr_factor=stripnet_lr_factor,
|
|
)
|
|
if t.learnable_temperature and self.loss_fn.logit_scale is not None:
|
|
param_groups[0]["params"].append(self.loss_fn.logit_scale)
|
|
self.optimizer = AdamW(param_groups, weight_decay=t.weight_decay)
|
|
|
|
lr_info = f"proj={t.learning_rate:.0e}"
|
|
if not self.models_common_cfg.baseline_mode:
|
|
lr_info += f" text={t.learning_rate * t.text_lr_factor:.0e}"
|
|
LOGGER.info(
|
|
"Optimizer: AdamW LR: %s warmup=%d epochs",
|
|
lr_info, self.pipeline_cfg.warmup_epochs,
|
|
)
|
|
|
|
# Scheduler — cosine with linear warmup (counted in optimizer steps).
|
|
self.steps_per_epoch = math.ceil(
|
|
len(self.train_loader) / self.hardware_cfg.grad_accum_steps,
|
|
)
|
|
total_steps = self.pipeline_cfg.epochs * self.steps_per_epoch
|
|
warmup_steps = self.pipeline_cfg.warmup_epochs * self.steps_per_epoch
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
|
|
self.scheduler = LambdaLR(
|
|
self.optimizer,
|
|
lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps),
|
|
last_epoch=-1,
|
|
)
|
|
self.scaler = GradScaler(enabled=self.hardware_cfg.use_amp)
|
|
|
|
def _restore_from_resume(self) -> None:
|
|
"""Restore optimizer/scheduler/loss state on resume."""
|
|
LOGGER.info("⚙️ Restore from resume...")
|
|
if self.resume_ckpt is None:
|
|
return
|
|
assert self.optimizer is not None and self.loss_fn is not None and self.scheduler is not None
|
|
if "optimizer_state" in self.resume_ckpt:
|
|
self.optimizer.load_state_dict(self.resume_ckpt["optimizer_state"])
|
|
LOGGER.info("Optimizer state restored")
|
|
if "loss_state" in self.resume_ckpt:
|
|
self.loss_fn.load_state_dict(self.resume_ckpt["loss_state"])
|
|
LOGGER.info("Loss state restored (tau=%.4f)", self.loss_fn.current_temperature)
|
|
# Set scheduler last_epoch so it resumes at the correct LR.
|
|
self.scheduler.last_epoch = self.start_epoch * self.steps_per_epoch
|
|
self.global_step = self.start_epoch * self.steps_per_epoch
|
|
LOGGER.info("Resuming from epoch %d", self.start_epoch)
|
|
|
|
def _setup_profiler(self) -> None:
|
|
"""Optional PyTorch profiler (only if start_epoch == 0)."""
|
|
LOGGER.info("⚙️ Setup profiler...")
|
|
if self.tracking_cfg.use_profiler and self.start_epoch == 0:
|
|
assert self.output_dir is not None
|
|
self.profiler = TrainingProfiler(
|
|
output_dir=self.output_dir,
|
|
n_warmup=self.tracking_cfg.profiler_warmup,
|
|
n_active=self.tracking_cfg.profiler_active,
|
|
)
|
|
self.profiler.start()
|
|
|
|
# ===================================================================
|
|
# Training loop
|
|
# ===================================================================
|
|
|
|
def _train_loop(self) -> None:
|
|
"""Main per-epoch loop."""
|
|
LOGGER.info(
|
|
"Starting training for %d epochs (from epoch %d)",
|
|
self.pipeline_cfg.epochs, self.start_epoch,
|
|
)
|
|
for epoch in range(self.start_epoch, self.pipeline_cfg.epochs):
|
|
self._maybe_dss_reembed(epoch)
|
|
epoch_start = time.time()
|
|
train_means = self._train_one_epoch(epoch)
|
|
elapsed = time.time() - epoch_start
|
|
self._log_epoch_summary(epoch, train_means, elapsed)
|
|
self._maybe_evaluate(epoch, train_means, elapsed)
|
|
self._save_checkpoint(epoch)
|
|
|
|
def _maybe_dss_reembed(self, epoch: int) -> None:
|
|
"""Refresh DSS query embeddings if the epoch is past warmup and on cadence."""
|
|
if not isinstance(self.batch_sampler, DynamicSimilaritySampler):
|
|
return
|
|
t = self.training_cfg
|
|
if epoch < t.dss_warmup_epochs:
|
|
return
|
|
if (epoch - t.dss_warmup_epochs) % t.dss_reembed_every != 0:
|
|
return
|
|
|
|
assert self.train_ds is not None and self.model is not None
|
|
query_embs: torch.Tensor | None = None
|
|
if self.emb_cache is not None:
|
|
query_embs = self.emb_cache.load(epoch)
|
|
if query_embs is None:
|
|
LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(self.train_ds), epoch)
|
|
t_embed = time.time()
|
|
query_embs = _embed_drone_queries(
|
|
self.model, self.train_ds, self.hardware_cfg.device,
|
|
batch_size=self.hardware_cfg.batch_size * self.hardware_cfg.grad_accum_steps,
|
|
num_workers=self.hardware_cfg.num_workers,
|
|
)
|
|
LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed)
|
|
if self.emb_cache is not None:
|
|
self.emb_cache.save(epoch, query_embs)
|
|
t_sampler = time.time()
|
|
self.batch_sampler.update_embeddings(query_embs)
|
|
LOGGER.info("DSS: sampler update_embeddings took %.2fs", time.time() - t_sampler)
|
|
|
|
def _train_one_epoch(self, epoch: int) -> dict[str, float]:
|
|
"""One epoch of training. Returns mean metrics (loss, temperature, gates)."""
|
|
assert self.model is not None and self.loss_fn is not None
|
|
assert self.optimizer is not None and self.scheduler is not None and self.scaler is not None
|
|
assert self.train_loader is not None and self.tracker is not None and self.csv_logger is not None
|
|
|
|
self.model.train()
|
|
if self.batch_sampler is not None:
|
|
self.batch_sampler.set_epoch(epoch)
|
|
|
|
agg: dict[str, float] = {}
|
|
n_batches = 0
|
|
accum = self.hardware_cfg.grad_accum_steps
|
|
device = self.hardware_cfg.device
|
|
baseline_mode = self.models_common_cfg.baseline_mode
|
|
|
|
pbar = tqdm(
|
|
self.train_loader,
|
|
desc=f" Epoch {epoch + 1}/{self.pipeline_cfg.epochs}",
|
|
unit="batch",
|
|
leave=False,
|
|
)
|
|
for batch in pbar:
|
|
if n_batches % accum == 0:
|
|
self.optimizer.zero_grad(set_to_none=True)
|
|
|
|
drone_img = batch["drone_img"].to(device, non_blocking=True)
|
|
sat_img = batch["sat_img"].to(device, non_blocking=True)
|
|
|
|
with autocast(device_type="cuda", enabled=self.hardware_cfg.use_amp):
|
|
if baseline_mode:
|
|
embeddings = self.model(drone_img=drone_img, sat_img=sat_img)
|
|
else:
|
|
embeddings = self.model(
|
|
drone_img=drone_img,
|
|
sat_img=sat_img,
|
|
caption_l1=batch["caption_l1"],
|
|
caption_l2=batch["caption_l2"],
|
|
caption_l3=batch["caption_l3"],
|
|
sat_caption_l1=batch["sat_caption_l1"],
|
|
sat_caption_l2=batch["sat_caption_l2"],
|
|
sat_caption_l3=batch["sat_caption_l3"]
|
|
)
|
|
|
|
queue_neg = self.neg_bank.get_queue() if self.neg_bank is not None else None
|
|
loss_kwargs: dict[str, Any] = {
|
|
"embeddings": embeddings,
|
|
"epoch": epoch,
|
|
"total_epochs": self.pipeline_cfg.epochs,
|
|
"queue_negatives": queue_neg,
|
|
}
|
|
if isinstance(self.loss_fn, WeightedInfoNCELoss):
|
|
loss_kwargs["positive_weights"] = batch["positive_weights"].to(
|
|
device, non_blocking=True,
|
|
)
|
|
loss_dict = self.loss_fn(**loss_kwargs)
|
|
|
|
raw_loss = float(loss_dict["total"].item())
|
|
total_loss = loss_dict["total"] / accum
|
|
self.scaler.scale(total_loss).backward()
|
|
|
|
if self.neg_bank is not None:
|
|
self.neg_bank.enqueue(embeddings["gallery"].detach())
|
|
|
|
is_accum_step = (n_batches + 1) % accum == 0 or (n_batches + 1) == len(self.train_loader)
|
|
if is_accum_step:
|
|
if self.training_cfg.grad_clip > 0:
|
|
self.scaler.unscale_(self.optimizer)
|
|
nn.utils.clip_grad_norm_(
|
|
self.model.trainable_parameters(),
|
|
max_norm=self.training_cfg.grad_clip,
|
|
)
|
|
if self.tracking_cfg.log_grad_norms and n_batches % (50 * accum) < accum:
|
|
grad_norms = compute_gradient_norms(self.model, self.loss_fn)
|
|
self.tracker.log_gradients(epoch, grad_norms, step=self.global_step)
|
|
if n_batches < accum:
|
|
log_gradient_summary(grad_norms)
|
|
self.scaler.step(self.optimizer)
|
|
self.scaler.update()
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
|
|
self.scheduler.step()
|
|
self.global_step += 1
|
|
|
|
step_metrics = {
|
|
"loss": raw_loss,
|
|
"temperature": float(loss_dict["temperature"].item()),
|
|
"gate_q": float(loss_dict["gate_q"].item()),
|
|
"gate_g": float(loss_dict["gate_g"].item()),
|
|
"lr": self.optimizer.param_groups[0]["lr"],
|
|
}
|
|
self.tracker.log_train(epoch, step_metrics, step=self.global_step)
|
|
self.csv_logger.log_batch(epoch, n_batches, self.global_step, step_metrics)
|
|
|
|
for key, val in loss_dict.items():
|
|
agg[key] = agg.get(key, 0.0) + float(val.item())
|
|
n_batches += 1
|
|
|
|
pbar.set_postfix(
|
|
loss=f"{raw_loss:.3f}",
|
|
tau=f"{step_metrics['temperature']:.4f}",
|
|
gq=f"{step_metrics['gate_q']:.3f}",
|
|
gg=f"{step_metrics['gate_g']:.3f}",
|
|
)
|
|
|
|
if self.profiler is not None:
|
|
self.profiler.step()
|
|
if self.profiler.is_done(n_batches):
|
|
self.profiler.export()
|
|
self.profiler = None
|
|
|
|
return {k: v / max(n_batches, 1) for k, v in agg.items()}
|
|
|
|
def _log_epoch_summary(
|
|
self, epoch: int, means: dict[str, float], elapsed: float,
|
|
) -> None:
|
|
"""Log epoch-level training summary + VRAM."""
|
|
assert self.optimizer is not None and self.tracker is not None
|
|
LOGGER.info(
|
|
"epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate_q=%.4f gate_g=%.4f",
|
|
epoch, elapsed,
|
|
self.optimizer.param_groups[0]["lr"],
|
|
means.get("total", 0.0),
|
|
means.get("temperature", 0.0),
|
|
means.get("gate_q", 1.0),
|
|
means.get("gate_g", 1.0),
|
|
)
|
|
if torch.cuda.is_available():
|
|
vram_gb = torch.cuda.max_memory_allocated() / 1e9
|
|
self.tracker.log_scalar("system/vram_peak_gb", vram_gb, step=self.global_step)
|
|
|
|
def _maybe_evaluate(
|
|
self, epoch: int, train_means: dict[str, float], elapsed: float,
|
|
) -> None:
|
|
"""Run eval if epoch % eval_every == 0 (or last epoch). Updates history + CSVs."""
|
|
assert self.model is not None and self.loss_fn is not None
|
|
assert self.train_eval_loader is not None and self.test_loader is not None
|
|
assert self.csv_logger is not None and self.tracker is not None and self.optimizer is not None
|
|
|
|
epoch_record: dict[str, Any] = {
|
|
"epoch": epoch,
|
|
"elapsed_seconds": elapsed,
|
|
"train": train_means,
|
|
}
|
|
|
|
is_last = epoch == self.pipeline_cfg.epochs - 1
|
|
run_eval = (epoch + 1) % self.pipeline_cfg.eval_every == 0 or is_last
|
|
|
|
train_recall: dict[str, float] = {}
|
|
if run_eval:
|
|
train_eval_batches = len(self.test_loader)
|
|
train_recall = evaluate(
|
|
self.model, self.train_eval_loader, self.hardware_cfg.device,
|
|
loss_fn=self.loss_fn, epoch=epoch, total_epochs=self.pipeline_cfg.epochs,
|
|
max_batches=train_eval_batches, desc="eval-train",
|
|
)
|
|
epoch_record["train_recall"] = train_recall
|
|
self.csv_logger.log_train_recall(epoch, train_recall)
|
|
self.tracker.log_train(
|
|
epoch,
|
|
{f"recall/{k}": v for k, v in train_recall.items() if k.startswith("r@")},
|
|
step=self.global_step,
|
|
)
|
|
|
|
# Train CSV row (combines means + recall if eval ran).
|
|
train_row: dict[str, float] = {**train_means}
|
|
if "total" in train_row:
|
|
train_row["train_loss"] = train_row.pop("total")
|
|
if train_recall:
|
|
for key in (
|
|
"r@1_q2g", "r@5_q2g", "r@10_q2g", "ap_q2g",
|
|
"r@1_g2q", "r@5_g2q", "r@10_g2q", "ap_g2q",
|
|
):
|
|
train_row[key] = train_recall.get(key, 0.0)
|
|
self.csv_logger.log_train(
|
|
epoch, train_row, self.optimizer.param_groups[0]["lr"], elapsed,
|
|
)
|
|
generate_plots(self.csv_logger.log_dir)
|
|
|
|
if train_recall:
|
|
LOGGER.info(
|
|
"train-recall epoch=%d q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f "
|
|
"g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f",
|
|
epoch,
|
|
train_recall.get("r@1_q2g", 0.0), train_recall.get("r@5_q2g", 0.0),
|
|
train_recall.get("r@10_q2g", 0.0), train_recall.get("ap_q2g", 0.0),
|
|
train_recall.get("r@1_g2q", 0.0), train_recall.get("r@5_g2q", 0.0),
|
|
train_recall.get("r@10_g2q", 0.0), train_recall.get("ap_g2q", 0.0),
|
|
train_recall.get("loss", 0.0),
|
|
)
|
|
|
|
val_metrics = evaluate(
|
|
self.model, self.test_loader, self.hardware_cfg.device,
|
|
loss_fn=self.loss_fn, epoch=epoch, total_epochs=self.pipeline_cfg.epochs,
|
|
desc="eval-val",
|
|
)
|
|
epoch_record["val"] = val_metrics
|
|
self.csv_logger.log_val(epoch, val_metrics)
|
|
generate_plots(self.csv_logger.log_dir)
|
|
self.tracker.log_val(epoch, val_metrics, step=self.global_step)
|
|
|
|
r1 = val_metrics.get("r@1_q2g", 0.0)
|
|
if r1 > self.best_r1:
|
|
self.best_r1 = r1
|
|
self.tracker.log_scalar(
|
|
"val/best_r@1_q2g", self.best_r1, step=self.global_step,
|
|
)
|
|
|
|
LOGGER.info(
|
|
"val epoch=%d q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f "
|
|
"g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f gate_q=%.4f",
|
|
epoch,
|
|
val_metrics.get("r@1_q2g", 0.0), val_metrics.get("r@5_q2g", 0.0),
|
|
val_metrics.get("r@10_q2g", 0.0), val_metrics.get("ap_q2g", 0.0),
|
|
val_metrics.get("r@1_g2q", 0.0), val_metrics.get("r@5_g2q", 0.0),
|
|
val_metrics.get("r@10_g2q", 0.0), val_metrics.get("ap_g2q", 0.0),
|
|
val_metrics.get("loss", 0.0), val_metrics.get("gate_q", 1.0),
|
|
)
|
|
|
|
self._maybe_run_gradcam(epoch)
|
|
|
|
self.history.append(epoch_record)
|
|
|
|
def _maybe_run_gradcam(self, epoch: int) -> None:
|
|
"""Optional Grad-CAM visualisation (only when gradcam_every divides epoch+1)."""
|
|
if not self.tracking_cfg.use_gradcam:
|
|
return
|
|
if (epoch + 1) % self.tracking_cfg.gradcam_every != 0:
|
|
return
|
|
from src.training.gradcam import generate_gradcam_samples
|
|
assert self.model is not None and self.test_loader is not None
|
|
assert self.output_dir is not None and self.tracker is not None
|
|
overlays = generate_gradcam_samples(
|
|
model=self.model,
|
|
dataloader=self.test_loader,
|
|
device=self.hardware_cfg.device,
|
|
output_dir=str(self.output_dir),
|
|
n_samples=self.tracking_cfg.gradcam_samples,
|
|
epoch=epoch,
|
|
)
|
|
for i, overlay in enumerate(overlays[:4]):
|
|
kind = "drone" if i % 2 == 0 else "sat"
|
|
self.tracker.log_image(
|
|
f"gradcam/{kind}_{i // 2}",
|
|
overlay,
|
|
step=self.global_step,
|
|
caption=f"Epoch {epoch} {kind} Grad-CAM",
|
|
)
|
|
|
|
def _save_checkpoint(self, epoch: int) -> None:
|
|
"""Save checkpoint atomically with backbone-specific arch flags."""
|
|
assert self.model is not None and self.optimizer is not None
|
|
assert self.loss_fn is not None and self.output_dir is not None
|
|
backbone = self.models_common_cfg.backbone
|
|
ckpt_obj: dict[str, Any] = {
|
|
"epoch": epoch,
|
|
"model_state": self.model.state_dict(),
|
|
"optimizer_state": self.optimizer.state_dict(),
|
|
"loss_state": self.loss_fn.state_dict(),
|
|
"baseline_mode": self.models_common_cfg.baseline_mode,
|
|
"backbone": backbone,
|
|
}
|
|
if isinstance(self.models_cfg, DINOv3ModelsConfig):
|
|
ckpt_obj["shared_encoder"] = self.models_cfg.shared_encoder
|
|
ckpt_obj["mona_bottleneck"] = self.models_cfg.mona_bottleneck
|
|
ckpt_obj["mona_last_n_blocks"] = self.models_cfg.mona_last_n_blocks
|
|
# StripNet: no extra arch flags saved here (params come from gin on resume).
|
|
atomic_save_torch(ckpt_obj, self.output_dir / f"ckpt_epoch{epoch:03d}.pt")
|
|
LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch)
|
|
|
|
# ===================================================================
|
|
# Final phase
|
|
# ===================================================================
|
|
|
|
def _final_evaluation(self) -> None:
|
|
"""Save history.json + eval_report.json + W&B summary."""
|
|
assert self.model is not None and self.loss_fn is not None
|
|
assert self.test_loader is not None and self.output_dir is not None
|
|
assert self.tracker is not None and self.full_config is not None
|
|
|
|
history_path = self.output_dir / "history.json"
|
|
with history_path.open("w", encoding="utf-8") as f:
|
|
json.dump(self.history, f, indent=2)
|
|
|
|
LOGGER.info("Running final evaluation...")
|
|
final_metrics = evaluate(
|
|
self.model, self.test_loader, self.hardware_cfg.device,
|
|
loss_fn=self.loss_fn,
|
|
epoch=self.pipeline_cfg.epochs - 1,
|
|
total_epochs=self.pipeline_cfg.epochs,
|
|
)
|
|
report = {
|
|
"config": self.full_config,
|
|
"metrics": final_metrics,
|
|
"history": self.history,
|
|
}
|
|
report_path = self.output_dir / "eval_report.json"
|
|
with report_path.open("w", encoding="utf-8") as f:
|
|
json.dump(report, f, indent=2)
|
|
|
|
self.tracker.log_summary({
|
|
"best_r@1_q2g": self.best_r1,
|
|
"final_r@1_q2g": final_metrics.get("r@1_q2g", 0.0),
|
|
"final_r@5_q2g": final_metrics.get("r@5_q2g", 0.0),
|
|
"final_r@10_q2g": final_metrics.get("r@10_q2g", 0.0),
|
|
"final_ap_q2g": final_metrics.get("ap_q2g", 0.0),
|
|
"final_r@1_g2q": final_metrics.get("r@1_g2q", 0.0),
|
|
"final_r@5_g2q": final_metrics.get("r@5_g2q", 0.0),
|
|
"final_r@10_g2q": final_metrics.get("r@10_g2q", 0.0),
|
|
"final_ap_g2q": final_metrics.get("ap_g2q", 0.0),
|
|
"final_gate_q": final_metrics.get("gate_q", 1.0),
|
|
"final_gate_g": final_metrics.get("gate_g", 1.0),
|
|
})
|
|
|
|
LOGGER.info("Training complete. Report: %s", report_path)
|
|
LOGGER.info(
|
|
"Final — q2g R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f "
|
|
"g2q R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f gate_q=%.4f gate_g=%.4f",
|
|
final_metrics.get("r@1_q2g", 0.0), final_metrics.get("r@5_q2g", 0.0),
|
|
final_metrics.get("r@10_q2g", 0.0), final_metrics.get("ap_q2g", 0.0),
|
|
final_metrics.get("r@1_g2q", 0.0), final_metrics.get("r@5_g2q", 0.0),
|
|
final_metrics.get("r@10_g2q", 0.0), final_metrics.get("ap_g2q", 0.0),
|
|
final_metrics.get("gate_q", 1.0), final_metrics.get("gate_g", 1.0),
|
|
)
|
|
|
|
def _cleanup(self) -> None:
|
|
"""Close profiler + tracker."""
|
|
if self.profiler is not None:
|
|
self.profiler.export()
|
|
self.profiler = None
|
|
if self.tracker is not None:
|
|
self.tracker.close()
|
|
|