More work on models!

evaluation/schema.py

@@ -296,7 +296,7 @@ def _git_commit() -> str:
 # ---------------------------------------------------------------------------
 def save_evaluation(ds: xr.Dataset, path: str) -> None:
     """Write an evaluation dataset to netCDF using the locked-down encodings."""
-    encoding = {name: _ENCODINGS[name] for name in ds.data_vars if name in _ENCODINGS}
+    encoding = {name: _ENCODINGS[name] for name in ds.data_vars if name in _ENCODINGS} # pyright: ignore
     ds.to_netcdf(path, engine=_NETCDF_ENGINE, encoding=encoding)
 
 

model/dnn_mod.py

@@ -36,6 +36,19 @@ from __future__ import absolute_import, division, print_function
 import bayesian_torch.layers as bayesian_layers
 from bayesian_torch.utils.util import get_rho
 
+# Canonical prior/posterior parameters for the DNN->BNN conversion. Training and
+# model loading MUST use the same values so a saved Bayesian state_dict maps onto
+# an identically-structured converted model.
+DEFAULT_BNN_PRIOR_PARAMETERS = {
+    "prior_mu": 0.0,
+    "prior_sigma": 1.0,
+    "posterior_mu_init": 0.0,
+    "posterior_rho_init": -3.0,
+    "type": "Reparameterization",
+    "moped_enable": False,
+    "moped_delta": 0.5,
+}
+
 # --------------------------------------------------------------------------------
 # Parameters used to define BNN layyers.
 #    bnn_prior_parameters = {

model/layers.py

@@ -44,30 +44,34 @@ class SplitCNVBlock(nn.Module):
 
         self.split_dim = split_dim
 
-        self.leftcnv_1 = SepCNVBlock(
-            in_channels // 2, mid_channels // 2, (3, 4, 3), droprate=drop_rate
+        # Build both branches once here. Previously these Sequentials were
+        # (re)assigned to ``self`` inside forward(), which registered
+        # ``leftblock``/``rightblock`` as submodules at runtime and duplicated
+        # their parameters into the state_dict. A freshly constructed model (no
+        # forward yet) lacked those keys, so a strict load_state_dict of a trained
+        # checkpoint failed with "Unexpected key(s)" at evaluation time.
+        self.leftblock = nn.Sequential(
+            SepCNVBlock(
+                in_channels // 2, mid_channels // 2, (3, 4, 3), droprate=drop_rate
+            ),
+            SepCNVBlock(
+                mid_channels // 2, out_channels // 2, (3, 4, 3), droprate=drop_rate
+            ),
         )
-        self.rightcnv_1 = SepCNVBlock(
-            in_channels // 2, mid_channels // 2, (3, 4, 3), droprate=drop_rate
-        )
-
-        self.leftcnv_2 = SepCNVBlock(
-            mid_channels // 2, out_channels // 2, (3, 4, 3), droprate=drop_rate
-        )
-        self.rightcnv_2 = SepCNVBlock(
-            mid_channels // 2, out_channels // 2, (3, 4, 3), droprate=drop_rate
+        self.rightblock = nn.Sequential(
+            SepCNVBlock(
+                in_channels // 2, mid_channels // 2, (3, 4, 3), droprate=drop_rate
+            ),
+            SepCNVBlock(
+                mid_channels // 2, out_channels // 2, (3, 4, 3), droprate=drop_rate
+            ),
         )
 
     def forward(self, x: Float[torch.Tensor, "N C D H W"]):
-        (left, right) = torch.tensor_split(x, 2, dim=self.split_dim)
-
-        self.leftblock = nn.Sequential(self.leftcnv_1, self.leftcnv_2)
-        self.rightblock = nn.Sequential(self.rightcnv_1, self.rightcnv_2)
-
+        left, right = torch.tensor_split(x, 2, dim=self.split_dim)
         left = self.leftblock(left)
         right = self.rightblock(right)
-        a = torch.cat((left, right), dim=self.split_dim)
-        return a
+        return torch.cat((left, right), dim=self.split_dim)
 
 
 class MidFlowBlock(nn.Module):

model/training.py

@@ -1,5 +1,4 @@
 import pathlib as pl
-import time
 from threading import Event
 from typing import Callable, Tuple, cast
 
@@ -10,6 +9,7 @@ import torch.nn as nn
 from torch.utils.data import DataLoader
 
 from model.dataset import ADNIDataset
+from util.control import check_control_events as _check_control_events
 from util.progress import ProgressTracker
 from util.ui_logger import PipelineLogger
 
@@ -62,18 +62,6 @@ def _batch_correct_and_total(
     return correct, total
 
 
-def _check_control_events(
-    stop_event: Event | None,
-    pause_event: Event | None,
-) -> None:
-    if stop_event is not None and stop_event.is_set():
-        raise InterruptedError("Pipeline execution stopped by user.")
-    while pause_event is not None and pause_event.is_set():
-        time.sleep(0.5)
-        if stop_event is not None and stop_event.is_set():
-            raise InterruptedError("Pipeline execution stopped by user while paused.")
-
-
 def test_model(
     model: nn.Module,
     test_loader: DataLoader[ADNIDataset],

tasks/__init__.py

@@ -8,43 +8,21 @@ Pipeline stages (see ai/ARCHITECTURE.md):
     1. load_data        -> implemented
     2. train_regular    -> implemented
     3. train_bayesian   -> implemented
-    4. evaluate_regular -> PLANNED (placeholder)
-    5. evaluate_bayesian-> PLANNED (placeholder)
-    6. evaluate_noisy   -> PLANNED (placeholder)
-    load_models         -> PLANNED (placeholder); loads saved .pt ensembles from
-                           disk into ``state`` so evaluation is decoupled from
-                           training (training deletes models from VRAM after saving).
+    4. evaluate_regular -> implemented
+    5. evaluate_bayesian-> implemented
+    6. evaluate_noisy   -> implemented
+    load_models         -> implemented; discovers saved .pt ensembles on disk and
+                           records their paths in ``state`` so evaluation is
+                           decoupled from training (training frees models from VRAM
+                           after saving). Evaluation loads one model at a time.
 """
 
-from typing import Any, Dict
-
-from util.progress import ProgressTracker
-from util.ui_logger import PipelineLogger
-
+from . import evaluate
 from . import load_data
+from . import load_models
 from . import train_bayesian
 from . import train_normal
 
-
-def _placeholder_task(planned: str):
-    """Build a no-op task for a pipeline stage that is not implemented yet.
-
-    It logs a clear warning and returns cleanly (advancing the pipeline) instead
-    of crashing, so the scenario wiring can be exercised before the real
-    evaluation logic lands.
-    """
-
-    def _task(
-        tracker: ProgressTracker,
-        logger: PipelineLogger,
-        config: Dict[str, Any],
-        state: Dict[str, Any],
-    ) -> None:
-        logger.info(f"[NOT IMPLEMENTED] {planned} — skipping (placeholder).")
-
-    return _task
-
-
 PIPELINE_TASKS = {
     "load_data": {
         "task_name": "Load Image and ADNIMERGE",
@@ -60,27 +38,19 @@ PIPELINE_TASKS = {
     },
     "load_models": {
         "task_name": "Load Saved Models",
-        "task_func": _placeholder_task(
-            "Load trained normal/Bayesian ensembles from work_dir into state"
-        ),
+        "task_func": load_models.load_models_task,
     },
     "evaluate_regular": {
         "task_name": "Evaluate Regular Models",
-        "task_func": _placeholder_task(
-            "Evaluate normal ensemble and save netCDF (step 4)"
-        ),
+        "task_func": evaluate.evaluate_normal_task,
     },
     "evaluate_bayesian": {
         "task_name": "Evaluate Bayesian Models",
-        "task_func": _placeholder_task(
-            "Evaluate Bayesian ensemble with MC uncertainty and save netCDF (step 5)"
-        ),
+        "task_func": evaluate.evaluate_bayesian_task,
     },
     "evaluate_noisy": {
         "task_name": "Evaluate Models on Noised Data",
-        "task_func": _placeholder_task(
-            "Evaluate both ensembles across Gaussian noise levels and save netCDF (step 6)"
-        ),
+        "task_func": evaluate.evaluate_noisy_task,
     },
 }
 

tasks/load_data.py

@@ -54,6 +54,9 @@ def load_data_task(
 
     ptids = list(zip(ptid_df["Image Data ID"].tolist(), ptid_df["PTID"].tolist()))
 
+    # Mapping used later by evaluation to attach patient ids to the sample axis.
+    state["image_to_ptid"] = {int(iid): str(pid) for iid, pid in ptids}
+
     # Split is grouped by PTID to prevent patient-level leakage across partitions.
     datasets = ds.divide_dataset_by_patient_id(
         dataset,

tasks/train_bayesian.py

@@ -1,7 +1,4 @@
-import gc
 import pathlib as pl
-import time
-from threading import Event
 from typing import Any, Dict
 
 import torch
@@ -11,7 +8,9 @@ from torch.utils.data import DataLoader
 import model.dataset as ds
 import model.training as tn
 from model.cnn import CNN3D
-from model.dnn_mod import dnn_to_bnn_mod, get_kl_loss
+from model.dnn_mod import DEFAULT_BNN_PRIOR_PARAMETERS, dnn_to_bnn_mod, get_kl_loss
+from util.control import check_control_events as _check_control_events
+from util.control import release_torch_memory as _release_torch_memory
 from util.progress import ProgressTracker
 from util.seeding import derive_seed, seed_everything
 from util.ui_logger import PipelineLogger
@@ -21,26 +20,6 @@ from util.ui_logger import PipelineLogger
 _BAYESIAN_SEED_STREAM = 1
 
 
-def _release_torch_memory(device: str) -> None:
-    gc.collect()
-    if device.startswith("cuda"):
-        torch.cuda.empty_cache()
-    elif device.startswith("mps") and hasattr(torch, "mps"):
-        torch.mps.empty_cache()
-
-
-def _check_control_events(
-    stop_event: Event | None,
-    pause_event: Event | None,
-) -> None:
-    if stop_event is not None and stop_event.is_set():
-        raise InterruptedError("Pipeline execution stopped by user.")
-    while pause_event is not None and pause_event.is_set():
-        time.sleep(0.5)
-        if stop_event is not None and stop_event.is_set():
-            raise InterruptedError("Pipeline execution stopped by user while paused.")
-
-
 def train_bayesian_task(
     track: ProgressTracker,
     log: PipelineLogger,
@@ -68,15 +47,9 @@ def train_bayesian_task(
     train_progress = track.get_sub_tracker("Training Progress")
     track.update(total=config["training"]["ensemble_size"], advance=0)
 
-    bnn_prior_parameters = {
-        "prior_mu": 0.0,
-        "prior_sigma": 1.0,
-        "posterior_mu_init": 0.0,
-        "posterior_rho_init": -3.0,
-        "type": "Reparameterization",
-        "moped_enable": False,
-        "moped_delta": 0.5,
-    }
+    # Shared with the model loader (evaluation) so saved state_dicts map onto an
+    # identically-converted model.
+    bnn_prior_parameters = DEFAULT_BNN_PRIOR_PARAMETERS
 
     base_seed = int(config["data"]["seed"])
 

tasks/train_normal.py

@@ -1,7 +1,4 @@
 import pathlib as pl
-import time
-import gc
-from threading import Event
 from typing import Any, Dict
 
 import torch
@@ -11,6 +8,8 @@ from torch.utils.data import DataLoader
 import model.dataset as ds
 import model.training as tn
 from model.cnn import CNN3D
+from util.control import check_control_events as _check_control_events
+from util.control import release_torch_memory as _release_torch_memory
 from util.progress import ProgressTracker
 from util.seeding import derive_seed, seed_everything
 from util.ui_logger import PipelineLogger
@@ -20,26 +19,6 @@ from util.ui_logger import PipelineLogger
 _NORMAL_SEED_STREAM = 0
 
 
-def _release_torch_memory(device: str) -> None:
-    gc.collect()
-    if device.startswith("cuda"):
-        torch.cuda.empty_cache()
-    elif device.startswith("mps") and hasattr(torch, "mps"):
-        torch.mps.empty_cache()
-
-
-def _check_control_events(
-    stop_event: Event | None,
-    pause_event: Event | None,
-) -> None:
-    if stop_event is not None and stop_event.is_set():
-        raise InterruptedError("Pipeline execution stopped by user.")
-    while pause_event is not None and pause_event.is_set():
-        time.sleep(0.5)
-        if stop_event is not None and stop_event.is_set():
-            raise InterruptedError("Pipeline execution stopped by user while paused.")
-
-
 def train_normal_task(
     track: ProgressTracker,
     log: PipelineLogger,