claude_refactor_v3: Add and passed test on splited gin-configs loading but without weigts
This commit is contained in:
@@ -37,6 +37,14 @@ def main() -> None:
|
|||||||
proj_dir = get_proj_dir()
|
proj_dir = get_proj_dir()
|
||||||
path2cfg = f"{proj_dir}in/config_files/" # per REQUIREMENTS_GIN_STYLE.md §5
|
path2cfg = f"{proj_dir}in/config_files/" # per REQUIREMENTS_GIN_STYLE.md §5
|
||||||
|
|
||||||
|
# -------------------------------------------------------
|
||||||
|
''' ONLY FOR DEBUG with launch.json config:
|
||||||
|
"args": ["main gtauav_balanced"] -> so need to extract
|
||||||
|
preset name "gtauav_balanced"
|
||||||
|
'''
|
||||||
|
preset_name = preset_name.split(' ')[1]
|
||||||
|
# -------------------------------------------------------
|
||||||
|
|
||||||
configs = load_all_configs(path2cfg, preset_name)
|
configs = load_all_configs(path2cfg, preset_name)
|
||||||
|
|
||||||
trainer = Trainer(
|
trainer = Trainer(
|
||||||
|
|||||||
@@ -192,23 +192,25 @@ class DINOv3ViT(nn.Module):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_pretrained(cls, path: str | Path) -> DINOv3ViT:
|
def from_pretrained(cls, path: str | Path) -> DINOv3ViT:
|
||||||
"""Load from .pth or .safetensors checkpoint."""
|
try:
|
||||||
model = cls()
|
"""Load from .pth or .safetensors checkpoint."""
|
||||||
path = Path(path)
|
model = cls()
|
||||||
LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
|
path = Path(path)
|
||||||
if path.suffix == ".safetensors":
|
LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
|
||||||
state = load_safetensors(str(path))
|
if path.suffix == ".safetensors":
|
||||||
else:
|
state = load_safetensors(str(path))
|
||||||
state = torch.load(str(path), map_location="cpu", weights_only=False)
|
else:
|
||||||
if "model" in state:
|
state = torch.load(str(path), map_location="cpu", weights_only=False)
|
||||||
state = state["model"]
|
if "model" in state:
|
||||||
elif "state_dict" in state:
|
state = state["model"]
|
||||||
state = state["state_dict"]
|
elif "state_dict" in state:
|
||||||
model.load_state_dict(state, strict=False)
|
state = state["state_dict"]
|
||||||
n_params = sum(p.numel() for p in model.parameters())
|
model.load_state_dict(state, strict=False)
|
||||||
LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
|
n_params = sum(p.numel() for p in model.parameters())
|
||||||
return model
|
LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
|
||||||
|
return model
|
||||||
|
except FileNotFoundError as e:
|
||||||
|
LOGGER.exception(msg=e.strerror)
|
||||||
|
|
||||||
# LRSCLIPTextEncoder removed — replaced by official DGTRS architecture
|
# LRSCLIPTextEncoder removed — replaced by official DGTRS architecture
|
||||||
# in src/models/dgtrs/model.py (DGTRSTextEncoder)
|
# in src/models/dgtrs/model.py (DGTRSTextEncoder)
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -230,7 +230,7 @@ class Trainer:
|
|||||||
def train(self) -> None:
|
def train(self) -> None:
|
||||||
"""Full pipeline: setup → build → train → evaluate → cleanup."""
|
"""Full pipeline: setup → build → train → evaluate → cleanup."""
|
||||||
self._validate_backbone()
|
self._validate_backbone()
|
||||||
clear_vram()
|
#! clear_vram()
|
||||||
set_seed(self.pipeline_cfg.seed)
|
set_seed(self.pipeline_cfg.seed)
|
||||||
self._setup_output_dir()
|
self._setup_output_dir()
|
||||||
self._setup_tracker()
|
self._setup_tracker()
|
||||||
@@ -256,6 +256,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _validate_backbone(self) -> None:
|
def _validate_backbone(self) -> None:
|
||||||
"""Reject unsupported backbones up front with a helpful message."""
|
"""Reject unsupported backbones up front with a helpful message."""
|
||||||
|
LOGGER.info("⚙️ Validate backbone")
|
||||||
backbone = self.models_common_cfg.backbone
|
backbone = self.models_common_cfg.backbone
|
||||||
if backbone not in _SUPPORTED_BACKBONES:
|
if backbone not in _SUPPORTED_BACKBONES:
|
||||||
raise NotImplementedError(
|
raise NotImplementedError(
|
||||||
@@ -271,6 +272,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _setup_output_dir(self) -> None:
|
def _setup_output_dir(self) -> None:
|
||||||
"""Create output_dir, save config.json, init csv_logger."""
|
"""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 = Path(self.pipeline_cfg.output_dir)
|
||||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
@@ -290,6 +292,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _setup_tracker(self) -> None:
|
def _setup_tracker(self) -> None:
|
||||||
"""W&B + TensorBoard tracker."""
|
"""W&B + TensorBoard tracker."""
|
||||||
|
LOGGER.info("⚙️ Setup tracker...")
|
||||||
assert self.output_dir is not None and self.full_config is not None
|
assert self.output_dir is not None and self.full_config is not None
|
||||||
self.tracker = ExperimentTracker(
|
self.tracker = ExperimentTracker(
|
||||||
output_dir=self.output_dir,
|
output_dir=self.output_dir,
|
||||||
@@ -303,6 +306,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _build_model(self) -> None:
|
def _build_model(self) -> None:
|
||||||
"""Build (or load) the encoder model based on the active backbone."""
|
"""Build (or load) the encoder model based on the active backbone."""
|
||||||
|
LOGGER.info("⚙️ Build loss...")
|
||||||
backbone = self.models_common_cfg.backbone
|
backbone = self.models_common_cfg.backbone
|
||||||
|
|
||||||
if self.pipeline_cfg.resume_from is not None:
|
if self.pipeline_cfg.resume_from is not None:
|
||||||
@@ -327,6 +331,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _build_model_from_resume(self, backbone: str) -> None:
|
def _build_model_from_resume(self, backbone: str) -> None:
|
||||||
"""Resume model from checkpoint. Sets self.model, self.resume_ckpt, self.start_epoch."""
|
"""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)
|
LOGGER.info("Resuming from %s", self.pipeline_cfg.resume_from)
|
||||||
# Both DINOv3 and StripNet go through AsymmetricEncoder.load_checkpoint.
|
# Both DINOv3 and StripNet go through AsymmetricEncoder.load_checkpoint.
|
||||||
# Note: load_checkpoint doesn't support StripNet — known existing limitation.
|
# Note: load_checkpoint doesn't support StripNet — known existing limitation.
|
||||||
@@ -348,6 +353,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _build_stripnet_model(self) -> nn.Module:
|
def _build_stripnet_model(self) -> nn.Module:
|
||||||
"""Construct AsymmetricEncoder configured for StripNet."""
|
"""Construct AsymmetricEncoder configured for StripNet."""
|
||||||
|
LOGGER.info("⚙️ Build StripNet model...")
|
||||||
assert isinstance(self.models_cfg, StripNetModelsConfig)
|
assert isinstance(self.models_cfg, StripNetModelsConfig)
|
||||||
m = self.models_cfg
|
m = self.models_cfg
|
||||||
# DINO paths passed but ignored at runtime when backbone='stripnet'.
|
# DINO paths passed but ignored at runtime when backbone='stripnet'.
|
||||||
@@ -370,6 +376,7 @@ class Trainer:
|
|||||||
).to(self.hardware_cfg.device)
|
).to(self.hardware_cfg.device)
|
||||||
|
|
||||||
def _build_dinov3_model(self) -> nn.Module:
|
def _build_dinov3_model(self) -> nn.Module:
|
||||||
|
LOGGER.info("⚙️ Build DINOv3 model...")
|
||||||
"""Construct AsymmetricEncoder configured for DINOv3."""
|
"""Construct AsymmetricEncoder configured for DINOv3."""
|
||||||
assert isinstance(self.models_cfg, DINOv3ModelsConfig)
|
assert isinstance(self.models_cfg, DINOv3ModelsConfig)
|
||||||
m = self.models_cfg
|
m = self.models_cfg
|
||||||
@@ -393,6 +400,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _configure_gradient_checkpointing(self) -> None:
|
def _configure_gradient_checkpointing(self) -> None:
|
||||||
"""Enable gradient checkpointing on encoders that support it."""
|
"""Enable gradient checkpointing on encoders that support it."""
|
||||||
|
LOGGER.info("⚙️ Configure gradient checkpointing...")
|
||||||
assert self.model is not None
|
assert self.model is not None
|
||||||
backbone = self.models_common_cfg.backbone
|
backbone = self.models_common_cfg.backbone
|
||||||
if not self.hardware_cfg.gradient_checkpointing:
|
if not self.hardware_cfg.gradient_checkpointing:
|
||||||
@@ -406,11 +414,11 @@ class Trainer:
|
|||||||
self.model.sat_encoder.set_gradient_checkpointing(True)
|
self.model.sat_encoder.set_gradient_checkpointing(True)
|
||||||
if self.model.text_encoder is not None:
|
if self.model.text_encoder is not None:
|
||||||
self.model.text_encoder.transformer.gradient_checkpointing = True
|
self.model.text_encoder.transformer.gradient_checkpointing = True
|
||||||
LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)")
|
LOGGER.info("✅ Gradient checkpointing enabled (DINOv3 + DGTRS)")
|
||||||
elif backbone == "stripnet":
|
elif backbone == "stripnet":
|
||||||
if self.model.text_encoder is not None:
|
if self.model.text_encoder is not None:
|
||||||
self.model.text_encoder.transformer.gradient_checkpointing = True
|
self.model.text_encoder.transformer.gradient_checkpointing = True
|
||||||
LOGGER.info("Gradient checkpointing enabled (DGTRS only; StripNet doesn't support it)")
|
LOGGER.info("✅ Gradient checkpointing enabled (DGTRS only; StripNet doesn't support it)")
|
||||||
|
|
||||||
def _log_model_summary(self) -> None:
|
def _log_model_summary(self) -> None:
|
||||||
"""Log trainable param count, save model_summary.txt, hook W&B."""
|
"""Log trainable param count, save model_summary.txt, hook W&B."""
|
||||||
@@ -428,6 +436,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _build_loss(self) -> None:
|
def _build_loss(self) -> None:
|
||||||
"""Build InfoNCELoss or WeightedInfoNCELoss based on training_cfg.loss_type."""
|
"""Build InfoNCELoss or WeightedInfoNCELoss based on training_cfg.loss_type."""
|
||||||
|
LOGGER.info("⚙️ Build loss...")
|
||||||
t = self.training_cfg
|
t = self.training_cfg
|
||||||
if t.loss_type == "symmetric":
|
if t.loss_type == "symmetric":
|
||||||
self.loss_fn = InfoNCELoss(
|
self.loss_fn = InfoNCELoss(
|
||||||
@@ -465,6 +474,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _build_neg_bank(self) -> None:
|
def _build_neg_bank(self) -> None:
|
||||||
"""Optional NegativeMemoryBank for hard-negative mining."""
|
"""Optional NegativeMemoryBank for hard-negative mining."""
|
||||||
|
LOGGER.info("⚙️ Build negative bank...")
|
||||||
assert self.model is not None
|
assert self.model is not None
|
||||||
if self.training_cfg.neg_bank_size > 0:
|
if self.training_cfg.neg_bank_size > 0:
|
||||||
self.neg_bank = NegativeMemoryBank(
|
self.neg_bank = NegativeMemoryBank(
|
||||||
@@ -478,6 +488,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _build_data_loaders(self) -> None:
|
def _build_data_loaders(self) -> None:
|
||||||
"""Build train/test/train_eval datasets, samplers, loaders."""
|
"""Build train/test/train_eval datasets, samplers, loaders."""
|
||||||
|
LOGGER.info("⚙️ Build dataloaders...")
|
||||||
drone_train_tf = get_drone_train_transform(image_size=256)
|
drone_train_tf = get_drone_train_transform(image_size=256)
|
||||||
sat_train_tf = get_satellite_train_transform(image_size=256)
|
sat_train_tf = get_satellite_train_transform(image_size=256)
|
||||||
eval_tf = get_dino_transform(image_size=256)
|
eval_tf = get_dino_transform(image_size=256)
|
||||||
@@ -594,6 +605,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _build_optimizer_and_scheduler(self) -> None:
|
def _build_optimizer_and_scheduler(self) -> None:
|
||||||
"""Build AdamW with per-group LR + cosine-warmup scheduler + GradScaler."""
|
"""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
|
assert self.model is not None and self.loss_fn is not None and self.train_loader is not None
|
||||||
t = self.training_cfg
|
t = self.training_cfg
|
||||||
|
|
||||||
@@ -637,6 +649,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _restore_from_resume(self) -> None:
|
def _restore_from_resume(self) -> None:
|
||||||
"""Restore optimizer/scheduler/loss state on resume."""
|
"""Restore optimizer/scheduler/loss state on resume."""
|
||||||
|
LOGGER.info("⚙️ Restore from resume...")
|
||||||
if self.resume_ckpt is None:
|
if self.resume_ckpt is None:
|
||||||
return
|
return
|
||||||
assert self.optimizer is not None and self.loss_fn is not None and self.scheduler is not None
|
assert self.optimizer is not None and self.loss_fn is not None and self.scheduler is not None
|
||||||
@@ -653,6 +666,7 @@ class Trainer:
|
|||||||
|
|
||||||
def _setup_profiler(self) -> None:
|
def _setup_profiler(self) -> None:
|
||||||
"""Optional PyTorch profiler (only if start_epoch == 0)."""
|
"""Optional PyTorch profiler (only if start_epoch == 0)."""
|
||||||
|
LOGGER.info("⚙️ Setup profiler...")
|
||||||
if self.tracking_cfg.use_profiler and self.start_epoch == 0:
|
if self.tracking_cfg.use_profiler and self.start_epoch == 0:
|
||||||
assert self.output_dir is not None
|
assert self.output_dir is not None
|
||||||
self.profiler = TrainingProfiler(
|
self.profiler = TrainingProfiler(
|
||||||
|
|||||||
294
tests/conftest.py
Normal file
294
tests/conftest.py
Normal file
@@ -0,0 +1,294 @@
|
|||||||
|
# """Shared pytest fixtures + reporter hooks for caption-test test suite.
|
||||||
|
|
||||||
|
# Provides:
|
||||||
|
# - proj_dir, path2cfg: real repository paths for integration-style tests
|
||||||
|
# that load actual gin presets from in/config_files/.
|
||||||
|
# - clear_gin: autouse fixture that wipes gin global state before every test
|
||||||
|
# (gin keeps bindings in module-level singleton; tests must not leak).
|
||||||
|
|
||||||
|
# Reporter hooks print "✅/❌ <docstring>" next to each test's PASSED/FAILED
|
||||||
|
# line in -v mode.
|
||||||
|
|
||||||
|
# Preset name lists live in a separate module (`tests/_presets.py`) so test
|
||||||
|
# files can import them with a plain `import _presets` — no relative imports,
|
||||||
|
# no need for tests/ to be a package.
|
||||||
|
# """
|
||||||
|
|
||||||
|
# from __future__ import annotations
|
||||||
|
|
||||||
|
# from pathlib import Path
|
||||||
|
|
||||||
|
# import gin
|
||||||
|
# import pytest
|
||||||
|
|
||||||
|
|
||||||
|
# DINOV3_PRESETS = (
|
||||||
|
# "gtauav_balanced",
|
||||||
|
# "gtauav_balanced_asym",
|
||||||
|
# "gtauav_baseline",
|
||||||
|
# "gtauav_baseline_asym",
|
||||||
|
# "gtauav_image_heavy",
|
||||||
|
# "gtauav_text_heavy",
|
||||||
|
# )
|
||||||
|
|
||||||
|
# STRIPNET_PRESETS = (
|
||||||
|
# "gtauav_balanced_stripnet",
|
||||||
|
# "gtauav_balanced_stripnet_unfrozen",
|
||||||
|
# "gtauav_baseline_stripnet",
|
||||||
|
# "gtauav_baseline_stripnet_unfrozen",
|
||||||
|
# )
|
||||||
|
|
||||||
|
# ALL_TRAINING_PRESETS = DINOV3_PRESETS + STRIPNET_PRESETS
|
||||||
|
|
||||||
|
# # --- gin hygiene -----------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
# @pytest.fixture(autouse=True)
|
||||||
|
# def clear_gin():
|
||||||
|
# """Wipe gin's global binding state before AND after each test.
|
||||||
|
|
||||||
|
# Gin keeps bindings in module-level singletons; without this fixture, a
|
||||||
|
# test that loads a config (or even just calls gin.parse_config_file in a
|
||||||
|
# helper) leaks bindings into the next test, leading to flaky failures
|
||||||
|
# like 'Unknown configurable' or wrong field values.
|
||||||
|
# """
|
||||||
|
# gin.clear_config()
|
||||||
|
# yield
|
||||||
|
# gin.clear_config()
|
||||||
|
|
||||||
|
|
||||||
|
# # --- real-repo paths -------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
# @pytest.fixture
|
||||||
|
# def proj_dir() -> Path:
|
||||||
|
# """Path to repository root (the directory that contains src/, tests/, in/)."""
|
||||||
|
# return Path(__file__).resolve().parent.parent
|
||||||
|
|
||||||
|
|
||||||
|
# @pytest.fixture
|
||||||
|
# def path2cfg(proj_dir: Path) -> str:
|
||||||
|
# """Trailing-slashed path to in/config_files/, matching `src/main.py`.
|
||||||
|
|
||||||
|
# Per REQUIREMENTS_GIN_STYLE.md §5, src/main.py builds this path as
|
||||||
|
# `f"{proj_dir}in/config_files/"`. Tests that exercise the real repo
|
||||||
|
# layout should use this fixture verbatim instead of constructing it
|
||||||
|
# independently.
|
||||||
|
# """
|
||||||
|
# return f"{proj_dir}/in/config_files/"
|
||||||
|
|
||||||
|
|
||||||
|
# # --- reporter hooks --------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
# def _docstring_summary(item: pytest.Item) -> str | None:
|
||||||
|
# """Return the first non-empty line of a test's docstring, or None."""
|
||||||
|
# func = getattr(item, "function", None) or getattr(item, "obj", None)
|
||||||
|
# if func is None or not getattr(func, "__doc__", None):
|
||||||
|
# return None
|
||||||
|
# for line in func.__doc__.strip().splitlines():
|
||||||
|
# stripped = line.strip()
|
||||||
|
# if stripped:
|
||||||
|
# return stripped
|
||||||
|
# return None
|
||||||
|
|
||||||
|
|
||||||
|
# # Cache nodeid → docstring summary, populated at collection time so the
|
||||||
|
# # logreport hook can look them up without re-introspecting the test function.
|
||||||
|
# _LAST_SEEN_SUMMARY: dict[str, str] = {}
|
||||||
|
|
||||||
|
|
||||||
|
# def pytest_collection_modifyitems(
|
||||||
|
# config: pytest.Config,
|
||||||
|
# items: list[pytest.Item],
|
||||||
|
# ) -> None:
|
||||||
|
# """Cache each item's docstring summary for later use by the status hook."""
|
||||||
|
# for item in items:
|
||||||
|
# summary = _docstring_summary(item)
|
||||||
|
# if summary:
|
||||||
|
# _LAST_SEEN_SUMMARY[item.nodeid] = summary
|
||||||
|
|
||||||
|
|
||||||
|
# def pytest_runtest_logreport(report: pytest.TestReport) -> None:
|
||||||
|
# """Print '✅/❌/⏭️ <docstring>' after each test's `call` phase finishes.
|
||||||
|
|
||||||
|
# Pytest emits 3 reports per test (setup → call → teardown). We hook the
|
||||||
|
# `call` phase — the one where pass/fail is actually decided — and write
|
||||||
|
# the icon + docstring summary to stdout, where it appears next to
|
||||||
|
# pytest's own PASSED/FAILED line.
|
||||||
|
# """
|
||||||
|
# if report.when != "call":
|
||||||
|
# return
|
||||||
|
|
||||||
|
# summary = _LAST_SEEN_SUMMARY.get(report.nodeid, "(no docstring)")
|
||||||
|
# if report.passed:
|
||||||
|
# icon = "✅"
|
||||||
|
# elif report.failed:
|
||||||
|
# icon = "❌"
|
||||||
|
# else:
|
||||||
|
# icon = "⏭️"
|
||||||
|
# print(f" {icon} {summary}")
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
"""Shared pytest fixtures + reporter hooks for caption-test test suite.
|
||||||
|
|
||||||
|
Provides:
|
||||||
|
- proj_dir, path2cfg: real repository paths for integration-style tests
|
||||||
|
that load actual gin presets from in/config_files/.
|
||||||
|
- clear_gin: autouse fixture that wipes gin global state before every test
|
||||||
|
(gin keeps bindings in module-level singleton; tests must not leak).
|
||||||
|
|
||||||
|
Reporter hooks print "✅/❌ <docstring>" next to each test's PASSED/FAILED
|
||||||
|
line in -v mode.
|
||||||
|
|
||||||
|
Preset name lists live in a separate module (`tests/_presets.py`) so test
|
||||||
|
files can import them with a plain `import _presets` — no relative imports,
|
||||||
|
no need for tests/ to be a package.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import gin
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
DINOV3_PRESETS = (
|
||||||
|
"gtauav_balanced",
|
||||||
|
"gtauav_balanced_asym",
|
||||||
|
"gtauav_baseline",
|
||||||
|
"gtauav_baseline_asym",
|
||||||
|
"gtauav_image_heavy",
|
||||||
|
"gtauav_text_heavy",
|
||||||
|
)
|
||||||
|
|
||||||
|
STRIPNET_PRESETS = (
|
||||||
|
"gtauav_balanced_stripnet",
|
||||||
|
"gtauav_balanced_stripnet_unfrozen",
|
||||||
|
"gtauav_baseline_stripnet",
|
||||||
|
"gtauav_baseline_stripnet_unfrozen",
|
||||||
|
)
|
||||||
|
|
||||||
|
ALL_TRAINING_PRESETS = DINOV3_PRESETS + STRIPNET_PRESETS
|
||||||
|
|
||||||
|
# --- gin hygiene -----------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(autouse=True)
|
||||||
|
def clear_gin():
|
||||||
|
"""Wipe gin's global binding state before AND after each test.
|
||||||
|
|
||||||
|
Gin keeps bindings in module-level singletons; without this fixture, a
|
||||||
|
test that loads a config (or even just calls gin.parse_config_file in a
|
||||||
|
helper) leaks bindings into the next test, leading to flaky failures
|
||||||
|
like 'Unknown configurable' or wrong field values.
|
||||||
|
"""
|
||||||
|
gin.clear_config()
|
||||||
|
yield
|
||||||
|
gin.clear_config()
|
||||||
|
|
||||||
|
|
||||||
|
# --- real-repo paths -------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def proj_dir() -> Path:
|
||||||
|
"""Path to repository root (the directory that contains src/, tests/, in/)."""
|
||||||
|
return Path(__file__).resolve().parent.parent
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def path2cfg(proj_dir: Path) -> str:
|
||||||
|
"""Trailing-slashed path to in/config_files/, matching `src/main.py`.
|
||||||
|
|
||||||
|
Per REQUIREMENTS_GIN_STYLE.md §5, src/main.py builds this path as
|
||||||
|
`f"{proj_dir}in/config_files/"`. Tests that exercise the real repo
|
||||||
|
layout should use this fixture verbatim instead of constructing it
|
||||||
|
independently.
|
||||||
|
"""
|
||||||
|
return f"{proj_dir}/in/config_files/"
|
||||||
|
|
||||||
|
|
||||||
|
# --- reporter hooks --------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def _docstring_summary(item: pytest.Item) -> str | None:
|
||||||
|
"""Return the first non-empty line of a test's docstring, or None."""
|
||||||
|
func = getattr(item, "function", None) or getattr(item, "obj", None)
|
||||||
|
if func is None or not getattr(func, "__doc__", None):
|
||||||
|
return None
|
||||||
|
for line in func.__doc__.strip().splitlines():
|
||||||
|
stripped = line.strip()
|
||||||
|
if stripped:
|
||||||
|
return stripped
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
# Cache nodeid → docstring summary, populated at collection time so the
|
||||||
|
# logreport hook can look them up without re-introspecting the test function.
|
||||||
|
_LAST_SEEN_SUMMARY: dict[str, str] = {}
|
||||||
|
|
||||||
|
|
||||||
|
def pytest_collection_modifyitems(
|
||||||
|
config: pytest.Config,
|
||||||
|
items: list[pytest.Item],
|
||||||
|
) -> None:
|
||||||
|
"""Cache each item's docstring summary for later use by the status hook."""
|
||||||
|
for item in items:
|
||||||
|
summary = _docstring_summary(item)
|
||||||
|
if summary:
|
||||||
|
_LAST_SEEN_SUMMARY[item.nodeid] = summary
|
||||||
|
|
||||||
|
|
||||||
|
def pytest_runtest_logreport(report: pytest.TestReport) -> None:
|
||||||
|
"""Print test results with parametrized tests grouped under one header.
|
||||||
|
|
||||||
|
Non-parametrized test:
|
||||||
|
✅ <docstring summary>
|
||||||
|
|
||||||
|
Parametrized test (first occurrence of the group):
|
||||||
|
<docstring summary>
|
||||||
|
✅ <param>
|
||||||
|
|
||||||
|
Parametrized test (subsequent occurrences):
|
||||||
|
✅ <param>
|
||||||
|
|
||||||
|
Pytest emits 3 reports per test (setup → call → teardown). We hook the
|
||||||
|
`call` phase — the one where pass/fail is actually decided.
|
||||||
|
"""
|
||||||
|
if report.when != "call":
|
||||||
|
return
|
||||||
|
|
||||||
|
summary = _LAST_SEEN_SUMMARY.get(report.nodeid, "(no docstring)")
|
||||||
|
if report.passed:
|
||||||
|
icon = "✅"
|
||||||
|
elif report.failed:
|
||||||
|
icon = "❌"
|
||||||
|
else:
|
||||||
|
icon = "⏭️"
|
||||||
|
|
||||||
|
# Detect parametrization: pytest encodes params as `nodeid[param1-param2-...]`.
|
||||||
|
if "[" in report.nodeid and report.nodeid.endswith("]"):
|
||||||
|
base_id, _, param_part = report.nodeid.partition("[")
|
||||||
|
param_label = param_part[:-1] or "empty" # strip trailing ']'
|
||||||
|
|
||||||
|
# Print docstring header only on first encounter of this group.
|
||||||
|
# Leading "\n" separates the header from pytest's progress dot.
|
||||||
|
if base_id not in _PRINTED_GROUP_HEADERS:
|
||||||
|
print(f"\n{summary}")
|
||||||
|
_PRINTED_GROUP_HEADERS.add(base_id)
|
||||||
|
print(f" {icon} {param_label}")
|
||||||
|
else:
|
||||||
|
print(f" {icon} {summary}")
|
||||||
|
|
||||||
|
|
||||||
|
# Tracks which parametrized test groups have already had their docstring
|
||||||
|
# header printed. Reset implicitly at the start of each pytest run because
|
||||||
|
# Python re-imports conftest.py.
|
||||||
|
_PRINTED_GROUP_HEADERS: set[str] = set()
|
||||||
187
tests/test_trainer.py
Normal file
187
tests/test_trainer.py
Normal file
@@ -0,0 +1,187 @@
|
|||||||
|
"""Tests for src.training.trainer_new.Trainer.
|
||||||
|
|
||||||
|
Scope: __init__ behaviour, backbone validation, ModelsConfig type union.
|
||||||
|
Out of scope: actual training (requires GPU + datasets + model checkpoints).
|
||||||
|
|
||||||
|
The Trainer class is designed to defer all heavy lifting (CUDA, model
|
||||||
|
construction, dataset loading) to .train(); __init__ just stores the 6 cfg
|
||||||
|
objects and zeros out runtime state. This makes it cheap to test.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import get_args
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from src.conf.config_loader import load_all_configs
|
||||||
|
from src.conf.models_dinov3_conf import DINOv3ModelsConfig
|
||||||
|
from src.conf.models_stripnet_conf import StripNetModelsConfig
|
||||||
|
from src.training.trainer_new import (
|
||||||
|
ModelsConfig,
|
||||||
|
Trainer,
|
||||||
|
_SUPPORTED_BACKBONES,
|
||||||
|
)
|
||||||
|
|
||||||
|
from conftest import DINOV3_PRESETS, STRIPNET_PRESETS
|
||||||
|
|
||||||
|
|
||||||
|
# -- module-level constants ------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def test_supported_backbones_is_frozenset() -> None:
|
||||||
|
"""_SUPPORTED_BACKBONES must be a frozenset (immutable, hashable)."""
|
||||||
|
assert isinstance(_SUPPORTED_BACKBONES, frozenset)
|
||||||
|
|
||||||
|
|
||||||
|
def test_supported_backbones_contents() -> None:
|
||||||
|
"""Exactly dinov3 and stripnet are supported in the current refactor.
|
||||||
|
|
||||||
|
Sofia (v1/v71) is intentionally absent — see Trainer._validate_backbone
|
||||||
|
for the rationale and the steps to add it later.
|
||||||
|
"""
|
||||||
|
assert _SUPPORTED_BACKBONES == frozenset({"dinov3", "stripnet"})
|
||||||
|
|
||||||
|
|
||||||
|
def test_models_config_union_contents() -> None:
|
||||||
|
"""ModelsConfig union mirrors _SUPPORTED_BACKBONES (dinov3 | stripnet)."""
|
||||||
|
union_members = set(get_args(ModelsConfig))
|
||||||
|
assert union_members == {DINOv3ModelsConfig, StripNetModelsConfig}
|
||||||
|
|
||||||
|
|
||||||
|
# -- _validate_backbone --------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("backbone", ["dinov3", "stripnet"])
|
||||||
|
def test_validate_backbone_accepts_supported(
|
||||||
|
path2cfg: str, backbone: str,
|
||||||
|
) -> None:
|
||||||
|
"""Supported backbones pass _validate_backbone silently.
|
||||||
|
|
||||||
|
We use a real preset to build a valid Trainer — this also exercises
|
||||||
|
the load_all_configs → Trainer(...) integration.
|
||||||
|
"""
|
||||||
|
preset = "gtauav_balanced" if backbone == "dinov3" else "gtauav_balanced_stripnet"
|
||||||
|
cfgs = load_all_configs(path2cfg, preset)
|
||||||
|
trainer = Trainer(
|
||||||
|
pipeline_cfg=cfgs["pipeline"],
|
||||||
|
hardware_cfg=cfgs["hardware"],
|
||||||
|
training_cfg=cfgs["training"],
|
||||||
|
tracking_cfg=cfgs["tracking"],
|
||||||
|
models_common_cfg=cfgs["models_common"],
|
||||||
|
models_cfg=cfgs["models"],
|
||||||
|
)
|
||||||
|
# Must not raise.
|
||||||
|
trainer._validate_backbone()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("bad_backbone", ["sofia_v1", "sofia_v71", "mistral_42b", ""])
|
||||||
|
def test_validate_backbone_rejects_unsupported(
|
||||||
|
path2cfg: str, bad_backbone: str,
|
||||||
|
) -> None:
|
||||||
|
"""Unsupported backbones (incl. sofia) raise NotImplementedError, not ImportError.
|
||||||
|
|
||||||
|
The user must get a clear, actionable message — not a stack trace from
|
||||||
|
a missing module.
|
||||||
|
"""
|
||||||
|
cfgs = load_all_configs(path2cfg, "gtauav_balanced")
|
||||||
|
trainer = Trainer(
|
||||||
|
pipeline_cfg=cfgs["pipeline"],
|
||||||
|
hardware_cfg=cfgs["hardware"],
|
||||||
|
training_cfg=cfgs["training"],
|
||||||
|
tracking_cfg=cfgs["tracking"],
|
||||||
|
models_common_cfg=cfgs["models_common"],
|
||||||
|
models_cfg=cfgs["models"],
|
||||||
|
)
|
||||||
|
# Tamper with backbone — simulate what would happen if config_loader
|
||||||
|
# were extended to accept sofia presets.
|
||||||
|
trainer.models_common_cfg.backbone = bad_backbone
|
||||||
|
|
||||||
|
with pytest.raises(NotImplementedError) as excinfo:
|
||||||
|
trainer._validate_backbone()
|
||||||
|
|
||||||
|
# Error message must mention the offending backbone name and what's supported.
|
||||||
|
msg = str(excinfo.value)
|
||||||
|
assert bad_backbone in msg or repr(bad_backbone) in msg
|
||||||
|
assert "dinov3" in msg
|
||||||
|
assert "stripnet" in msg
|
||||||
|
|
||||||
|
|
||||||
|
# -- Trainer.__init__ smoke tests ------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("preset_name", DINOV3_PRESETS + STRIPNET_PRESETS)
|
||||||
|
def test_trainer_init_with_real_preset(path2cfg: str, preset_name: str) -> None:
|
||||||
|
"""Trainer(...) instantiates from every real preset's loaded cfgs.
|
||||||
|
|
||||||
|
Heavy work (CUDA, model build, dataset open) is deferred to .train();
|
||||||
|
__init__ only stores cfgs and zeros runtime state, so this is cheap and
|
||||||
|
GPU-free.
|
||||||
|
"""
|
||||||
|
cfgs = load_all_configs(path2cfg, preset_name)
|
||||||
|
|
||||||
|
trainer = Trainer(
|
||||||
|
pipeline_cfg=cfgs["pipeline"],
|
||||||
|
hardware_cfg=cfgs["hardware"],
|
||||||
|
training_cfg=cfgs["training"],
|
||||||
|
tracking_cfg=cfgs["tracking"],
|
||||||
|
models_common_cfg=cfgs["models_common"],
|
||||||
|
models_cfg=cfgs["models"],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Cfgs are stored as-is.
|
||||||
|
assert trainer.pipeline_cfg is cfgs["pipeline"]
|
||||||
|
assert trainer.hardware_cfg is cfgs["hardware"]
|
||||||
|
assert trainer.training_cfg is cfgs["training"]
|
||||||
|
assert trainer.tracking_cfg is cfgs["tracking"]
|
||||||
|
assert trainer.models_common_cfg is cfgs["models_common"]
|
||||||
|
assert trainer.models_cfg is cfgs["models"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_trainer_init_zeros_runtime_state(path2cfg: str) -> None:
|
||||||
|
"""All runtime fields are None / 0 / [] before .train() is called."""
|
||||||
|
cfgs = load_all_configs(path2cfg, "gtauav_balanced")
|
||||||
|
trainer = Trainer(
|
||||||
|
pipeline_cfg=cfgs["pipeline"],
|
||||||
|
hardware_cfg=cfgs["hardware"],
|
||||||
|
training_cfg=cfgs["training"],
|
||||||
|
tracking_cfg=cfgs["tracking"],
|
||||||
|
models_common_cfg=cfgs["models_common"],
|
||||||
|
models_cfg=cfgs["models"],
|
||||||
|
)
|
||||||
|
|
||||||
|
# None-typed runtime fields.
|
||||||
|
for attr in (
|
||||||
|
"output_dir", "full_config", "tracker", "csv_logger", "model",
|
||||||
|
"loss_fn", "neg_bank", "optimizer", "scheduler", "scaler",
|
||||||
|
"train_ds", "test_ds", "train_eval_ds",
|
||||||
|
"train_loader", "test_loader", "train_eval_loader",
|
||||||
|
"batch_sampler", "emb_cache", "profiler", "resume_ckpt",
|
||||||
|
):
|
||||||
|
assert getattr(trainer, attr) is None, (
|
||||||
|
f"trainer.{attr} should be None before .train(), "
|
||||||
|
f"got {type(getattr(trainer, attr)).__name__}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Counter / loop state initialized to identity values.
|
||||||
|
assert trainer.start_epoch == 0
|
||||||
|
assert trainer.global_step == 0
|
||||||
|
assert trainer.best_r1 == 0.0
|
||||||
|
assert trainer.history == []
|
||||||
|
assert trainer.steps_per_epoch == 0
|
||||||
|
|
||||||
|
|
||||||
|
# -- Trainer.train end-to-end signature ------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def test_trainer_train_method_exists_and_takes_no_args() -> None:
|
||||||
|
"""Trainer.train() takes only `self` — main.py calls trainer.train()."""
|
||||||
|
import inspect
|
||||||
|
|
||||||
|
sig = inspect.signature(Trainer.train)
|
||||||
|
params = [p for p in sig.parameters.values() if p.name != "self"]
|
||||||
|
assert params == [], (
|
||||||
|
f"Trainer.train() must take only self; got extra params: {params}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
Reference in New Issue
Block a user