claude_refactor_v3: Updated main (entry point), trainer_new (last version of train_gtauav), check: is extracted evluate() from train to evaluator.py correct in new context

src/eval/evaluator.py

@@ -6,12 +6,10 @@ Computes R@K and MRR for both q→g (drone→satellite) and g→q (satellite→d
 on the full satellite gallery. Multi-match: a query counts as a hit@K if ANY
 of its valid satellite matches (sat_candidates) appears in the top-K.
 
-Body transplanted from src/training/train_gtauav.py::_evaluate (pre-step-4b)
-with two changes:
-  1. Decorator @torch.no_grad() → @torch.inference_mode().
-  2. Type annotation `model: AsymmetricEncoder` → `model: nn.Module`
-     (any encoder with encode_query/encode_gallery + fusion_query.gate_value
-     and fusion_gallery.gate_value duck-typed attributes).
+Body transplanted byte-for-byte from src/training/train_gtauav.py::_evaluate
+in the main branch. The single difference is the type annotation
+`model: AsymmetricEncoder` → `model: nn.Module` (relaxed for duck-typing
+across encoder families); semantically identical to the main-branch version.
 
 Note: not to be confused with src/eval/evaluate.py (legacy v2 helper for
 UAV-VisLoc with a different signature). This module lives at
@@ -19,13 +17,13 @@ src/eval/evaluator.py and is the active evaluator for v3 GTA-UAV-LR.
 """
 
 import logging
-from typing import Any
 
 import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
+from src.models.asymmetric_encoder import AsymmetricEncoder
 from src.datasets.gtauav_dataset import (
     GTAUAVDataset,
     GTAUAVDroneQuery,
@@ -37,9 +35,9 @@ from src.datasets.gtauav_dataset import (
 LOGGER = logging.getLogger("caption_test.evaluator")
 
 
-@torch.inference_mode()
+@torch.no_grad()
 def evaluate(
-    model: nn.Module,
+    model: AsymmetricEncoder,
     loader: DataLoader,
     device: str,
     loss_fn: nn.Module | None = None,
@@ -57,8 +55,11 @@ def evaluate(
     satellite matches (pair_pos_sate_img_list ∪ pair_pos_semipos_sate_img_list)
     appears in the top-K.
 
+    `max_batches` subsamples the drone queries (not the gallery) — useful
+    for a quick train-side sanity check.
+
     Args:
-        model: Encoder with `encode_query(drone_img, l1, l2, l3, altitude=...)`
+        model: Encoder with `encode_query(drone_img, l1, l2, l3)`
                and `encode_gallery(sat_img, l1, l2, l3)`. Must expose
                `fusion_query.gate_value` and `fusion_gallery.gate_value`.
         loader: DataLoader over a GTAUAVDataset (used only to pull dataset
@@ -133,13 +134,9 @@ def evaluate(
         if max_batches is not None and i >= max_batches:
             break
         drone_img = batch["drone_img"].to(device, non_blocking=True)
-        altitude = batch.get("altitude")
-        if altitude is not None:
-            altitude = altitude.to(device, non_blocking=True)
         q = model.encode_query(
             drone_img,
             batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
-            altitude=altitude,
         )
         query_embs.append(q.cpu())
         query_valid_names.extend(batch["valid_sat_names"])

src/main.py

@@ -2,7 +2,7 @@
 
 Usage:
     python -m src.main gtauav_balanced
-    python -m src.main gtauav_balanced_sofia_v1
+    python -m src.main gtauav_balanced_stripnet
 """
 
 from __future__ import annotations
@@ -13,7 +13,7 @@ import sys
 import coloredlogs
 
 from src.conf.config_loader import load_all_configs
-from src.training.train_gtauav_old import train
+from src.training.trainer_new import Trainer
 from src.utils.path_utils import get_proj_dir
 
 logger = logging.getLogger("caption_test")
@@ -39,7 +39,7 @@ def main() -> None:
 
     configs = load_all_configs(path2cfg, preset_name)
 
-    train(
+    trainer = Trainer(
         pipeline_cfg=configs["pipeline"],
         hardware_cfg=configs["hardware"],
         training_cfg=configs["training"],
@@ -48,7 +48,10 @@ def main() -> None:
         models_cfg=configs["models"],
     )
 
+    trainer.train()
+
 
 if __name__ == "__main__":
     main()
 
+

src/training/trainer_new.py

@@ -3,13 +3,13 @@ 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 `run()` plus dedicated `_setup_*` / `_build_*` /
+orchestrating method `train()` plus dedicated `_setup_*` / `_build_*` /
 `_train_*` / `_evaluate_*` methods.
 
 Lifecycle:
-    Trainer(...) → run() → done.
+    Trainer(...) → train() → done.
 
-`run()` calls _build_* in dependency order, then _train_loop, then
+`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
@@ -80,7 +80,7 @@ _SUPPORTED_BACKBONES: frozenset[str] = frozenset({"dinov3", "stripnet"})
 
 
 def _build_param_groups(
-    model: nn.Module,
+    model: AsymmetricEncoder,
     lr: float,
     text_lr_factor: float,
     stripnet_backbone_lr_factor: float = 0.1,
@@ -130,7 +130,7 @@ def _cosine_warmup_schedule(warmup_steps: int, total_steps: int):
 
 
 def _embed_drone_queries(
-    model: nn.Module,
+    model: AsymmetricEncoder,
     train_ds: GTAUAVDataset,
     device: str,
     batch_size: int,
@@ -155,7 +155,7 @@ def _embed_drone_queries(
     )
     all_embs: list[torch.Tensor] = []
     with torch.inference_mode():
-        for batch in tqdm(loader, desc="dss-embed", unit="batch", leave=False):
+        for batch in tqdm(loader, desc="  dss-embed-queries", unit="batch", leave=False):
             drone_img = batch["drone_img"].to(device, non_blocking=True)
             altitude = batch.get("altitude")
             if altitude is not None:
@@ -178,7 +178,7 @@ class Trainer:
 
     All gin parameters arrive as 6 config objects; runtime state (model,
     optimizer, loaders, ...) is built lazily by _build_* methods and lives
-    on `self`. `run()` calls them in dependency order.
+    on `self`. `train()` calls them in dependency order.
 
     Backbones supported: 'dinov3', 'stripnet'.
     """
@@ -232,7 +232,7 @@ class Trainer:
     # Public entry point
     # ===================================================================
 
-    def run(self) -> None:
+    def train(self) -> None:
         """Full pipeline: setup → build → train → evaluate → cleanup."""
         self._validate_backbone()
         clear_vram()
@@ -1052,4 +1052,3 @@ class Trainer:
         if self.tracker is not None:
             self.tracker.close()
 
-