new files as well

model/factory.py

@@ -0,0 +1,65 @@
+"""Construction and loading of ensemble models for evaluation.
+
+Training saves each member's ``state_dict`` to ``model_N.pt``. Evaluation must
+rebuild an identically-structured module before ``load_state_dict``. For Bayesian
+members this means applying the DNN->BNN conversion (with the same prior params
+used at training time) *before* loading weights, otherwise the parameter names
+won't match.
+"""
+
+import pathlib as pl
+from typing import Any, Dict
+
+import torch
+from torch import nn
+
+from model.cnn import CNN3D
+from model.dnn_mod import DEFAULT_BNN_PRIOR_PARAMETERS, dnn_to_bnn_mod
+
+# Model-kind strings (mirror evaluation.schema.MODEL_KIND_*). Kept here too so
+# model/ has no dependency on the evaluation package.
+KIND_NORMAL = "normal"
+KIND_BAYESIAN = "bayesian"
+
+
+def build_model(config: Dict[str, Any], kind: str) -> CNN3D:
+    """Build an (untrained) CNN3D of the requested kind on CPU.
+
+    For ``kind == "bayesian"`` the conv/linear layers are converted to their
+    Bayesian counterparts using :data:`DEFAULT_BNN_PRIOR_PARAMETERS`.
+    """
+    data_cfg = config["data"]
+    train_cfg = config["training"]
+    model = CNN3D(
+        image_channels=data_cfg["image_channels"],
+        clin_data_channels=data_cfg["clin_data_channels"],
+        num_classes=data_cfg["num_classes"],
+        droprate=train_cfg["droprate"],
+    ).float()
+
+    if kind == KIND_BAYESIAN:
+        dnn_to_bnn_mod(model, DEFAULT_BNN_PRIOR_PARAMETERS)
+    elif kind != KIND_NORMAL:
+        raise ValueError(f"Unknown model kind {kind!r}; expected 'normal'/'bayesian'.")
+
+    return model
+
+
+def load_model(
+    path: pl.Path,
+    kind: str,
+    config: Dict[str, Any],
+    device: str,
+) -> nn.Module:
+    """Rebuild a model of ``kind`` and load ``path``'s weights onto ``device``.
+
+    The returned model is in ``eval()`` mode. Bayesian models remain stochastic
+    in eval mode (their reparameterization layers sample every forward pass),
+    which is exactly what Monte-Carlo evaluation relies on.
+    """
+    model = build_model(config, kind)
+    state_dict = torch.load(path, map_location="cpu")
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+    return model

tasks/evaluate.py

@@ -0,0 +1,332 @@
+"""Evaluation tasks (pipeline steps 4-6).
+
+All three tasks share one core, :func:`_evaluate_ensemble`, which loads each
+ensemble member one at a time, runs Monte-Carlo forward passes over the (fixed,
+unshuffled) test set at each requested noise level, and writes the results to a
+netCDF file in the locked-down evaluation schema (:mod:`evaluation.schema`).
+
+    step 4  evaluate_regular  -> normal ensemble, noise=[0.0], K=1
+    step 5  evaluate_bayesian -> Bayesian ensemble, noise=[0.0], K=mc_passes
+    step 6  evaluate_noisy    -> both ensembles across config noise_levels
+
+Determinism
+-----------
+The test loader is unshuffled, so sample order is stable and aligns with the
+netCDF ``sample`` axis. Added Gaussian noise is drawn from a generator seeded by
+``derive_seed(base_seed, noise_level_index, stream=_NOISE_SEED_STREAM)`` and
+re-seeded at the start of every MC pass, so:
+  * the noise perturbation for a given (sample, level) is identical across the K
+    MC passes (only the sampled BNN weights vary), and
+  * every model sees the exact same noised inputs at a given level (fair
+    comparison across the ensemble).
+"""
+
+import pathlib as pl
+from typing import Any, Dict, List, Sequence
+
+import numpy as np
+import toml
+import torch
+
+from evaluation.schema import (
+    CLASS_NAMES,
+    EvaluationAccumulator,
+    save_evaluation,
+)
+from model.factory import KIND_BAYESIAN, KIND_NORMAL, load_model
+from util.control import check_control_events, release_torch_memory
+from util.progress import ProgressTracker
+from util.seeding import derive_seed
+from util.ui_logger import PipelineLogger
+
+# RNG stream id for input noise (0=normal training, 1=Bayesian training).
+_NOISE_SEED_STREAM = 2
+
+# Subdirectory (under work_dir) where evaluation netCDF files are written.
+_EVAL_SUBDIR = "evaluations"
+
+# Fallbacks if a config predates the [evaluation] section.
+_DEFAULT_NOISE_LEVELS: List[float] = [0.0, 0.02, 0.05, 0.1, 0.2]
+_DEFAULT_MC_PASSES = 30
+
+
+def _events(state: Dict[str, Any]):
+    events = state.get("events")
+    stop_event = events.get("stop") if isinstance(events, dict) else None
+    pause_event = events.get("pause") if isinstance(events, dict) else None
+    return stop_event, pause_event
+
+
+def _collect_sample_metadata(loader, image_to_ptid: Dict[int, str]):
+    """Return ordered ``(image_ids, ptids, true_classes)`` for the test set.
+
+    The order matches iteration of the (unshuffled) loader and therefore the
+    order used when accumulating model outputs.
+    """
+    image_ids: List[int] = []
+    true_classes: List[str] = []
+    for _mri, _xls, targets, ids in loader:
+        for i in range(int(targets.shape[0])):
+            image_ids.append(int(ids[i]))
+            true_classes.append(CLASS_NAMES[int(targets[i].argmax())])
+    ptids = [str(image_to_ptid.get(iid, "unknown")) for iid in image_ids]
+    return image_ids, ptids, true_classes
+
+
+def _run_mc_passes(
+    model: torch.nn.Module,
+    loader,
+    device: str,
+    noise_level: float,
+    n_passes: int,
+    noise_seed: int,
+    n_classes: int,
+    parent_tracker: ProgressTracker,
+    stop_event,
+    pause_event,
+    label: str,
+) -> np.ndarray:
+    """Run ``n_passes`` forward passes over the test set, returning ``(K, S, C)``.
+
+    For a deterministic model use ``n_passes == 1``. Bayesian models sample fresh
+    weights each pass (their reparameterization layers stay stochastic in eval
+    mode), producing the Monte-Carlo ensemble the uncertainty measures need.
+    """
+    n_samples = len(loader.dataset)
+    mc = np.zeros((n_passes, n_samples, n_classes), dtype=np.float32)
+
+    pass_tracker = parent_tracker.get_sub_tracker(label)
+    pass_tracker.update(total=n_passes, advance=0)
+
+    with torch.no_grad():
+        for k in range(n_passes):
+            # Re-seed per pass so the noise realization is identical across MC
+            # passes (and across models) for a given (sample, noise_level).
+            noise_gen = (
+                torch.Generator(device=device).manual_seed(int(noise_seed))
+                if noise_level > 0.0
+                else None
+            )
+
+            batch_tracker = pass_tracker.get_sub_tracker("Batches")
+            batch_tracker.update(total=len(loader), advance=0)
+
+            idx = 0
+            for mri, xls, _targets, _ids in loader:
+                check_control_events(stop_event, pause_event)
+                mri = mri.to(device, non_blocking=True)
+                xls = xls.to(device, non_blocking=True)
+
+                if noise_gen is not None:
+                    noise = torch.randn(
+                        mri.shape, generator=noise_gen, device=device, dtype=mri.dtype
+                    )
+                    mri = mri + noise * noise_level
+
+                outputs = model((mri, xls))
+                bs = int(outputs.shape[0])
+                mc[k, idx : idx + bs] = outputs.detach().cpu().numpy()
+                idx += bs
+                batch_tracker.update(advance=1)
+
+            pass_tracker.update(advance=1)
+
+    return mc
+
+
+def _evaluate_ensemble(
+    *,
+    model_paths: Sequence[pl.Path],
+    kind: str,
+    config: Dict[str, Any],
+    state: Dict[str, Any],
+    log: PipelineLogger,
+    track: ProgressTracker,
+    noise_levels: Sequence[float],
+    n_passes: int,
+    out_path: pl.Path,
+) -> None:
+    """Evaluate one ensemble and write a schema-compliant netCDF to ``out_path``."""
+    device = config["training"]["device"]
+    n_classes = int(config["data"]["num_classes"])
+    base_seed = int(config["data"]["seed"])
+    stop_event, pause_event = _events(state)
+
+    test_loader = state["test_loader"]
+    image_to_ptid: Dict[int, str] = state.get("image_to_ptid", {})
+
+    image_ids, ptids, true_classes = _collect_sample_metadata(test_loader, image_to_ptid)
+    log.info(
+        f"Evaluating {len(model_paths)} {kind} model(s) on {len(image_ids)} samples, "
+        f"{len(noise_levels)} noise level(s), {n_passes} MC pass(es) each."
+    )
+
+    accumulator = EvaluationAccumulator(
+        image_ids, ptids, true_classes, noise_levels=list(noise_levels)
+    )
+
+    models_tracker = track.get_sub_tracker(f"{kind.capitalize()} models")
+    models_tracker.update(total=len(model_paths), advance=0)
+
+    for model_index, path in enumerate(model_paths):
+        check_control_events(stop_event, pause_event)
+        log.info(f"Loading {kind} model {model_index + 1}/{len(model_paths)}: {path.name}")
+        model = load_model(path, kind, config, device)
+
+        for level_index, noise in enumerate(noise_levels):
+            noise_seed = derive_seed(base_seed, level_index, stream=_NOISE_SEED_STREAM)
+            mc_probs = _run_mc_passes(
+                model=model,
+                loader=test_loader,
+                device=device,
+                noise_level=float(noise),
+                n_passes=n_passes,
+                noise_seed=noise_seed,
+                n_classes=n_classes,
+                parent_tracker=models_tracker,
+                stop_event=stop_event,
+                pause_event=pause_event,
+                label=f"model {model_index + 1} | sigma={noise}",
+            )
+            accumulator.add(
+                model_index=model_index,
+                model_kind=kind,
+                noise_level=float(noise),
+                mc_probs=mc_probs,
+            )
+
+        del model
+        release_torch_memory(device)
+        models_tracker.update(advance=1)
+
+    dataset = accumulator.to_dataset(
+        attrs={
+            "seed": base_seed,
+            "n_mc": n_passes,
+            "ensemble_kind": kind,
+            "config": _config_snapshot(config),
+        }
+    )
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    save_evaluation(dataset, str(out_path))
+    log.info(
+        f"Saved {kind} evaluation to {out_path} "
+        f"(models={dataset.sizes['model']}, samples={dataset.sizes['sample']}, "
+        f"noise_levels={dataset.sizes['noise_level']})."
+    )
+
+
+def _config_snapshot(config: Dict[str, Any]) -> str:
+    try:
+        return toml.dumps(config)
+    except Exception:
+        return repr(config)
+
+
+def _eval_cfg(config: Dict[str, Any]):
+    eval_cfg = config.get("evaluation", {})
+    noise_levels = list(eval_cfg.get("noise_levels", _DEFAULT_NOISE_LEVELS))
+    mc_passes = int(eval_cfg.get("mc_passes", _DEFAULT_MC_PASSES))
+    return noise_levels, mc_passes
+
+
+def _out_path(config: Dict[str, Any], name: str) -> pl.Path:
+    return pl.Path(config["work_dir"]) / _EVAL_SUBDIR / name
+
+
+# ---------------------------------------------------------------------------
+# Pipeline task entry points
+# ---------------------------------------------------------------------------
+def evaluate_normal_task(
+    track: ProgressTracker,
+    log: PipelineLogger,
+    config: Dict[str, Any],
+    state: Dict[str, Any],
+) -> None:
+    """Step 4: evaluate the normal ensemble on clean data (noise=0.0, K=1)."""
+    paths = state.get("normal_model_paths") or []
+    if not paths:
+        log.error("No normal models to evaluate (run training/load_models first).")
+        return
+    _evaluate_ensemble(
+        model_paths=paths,
+        kind=KIND_NORMAL,
+        config=config,
+        state=state,
+        log=log,
+        track=track,
+        noise_levels=[0.0],
+        n_passes=1,
+        out_path=_out_path(config, "normal.nc"),
+    )
+
+
+def evaluate_bayesian_task(
+    track: ProgressTracker,
+    log: PipelineLogger,
+    config: Dict[str, Any],
+    state: Dict[str, Any],
+) -> None:
+    """Step 5: evaluate the Bayesian ensemble on clean data with MC uncertainty."""
+    paths = state.get("bayesian_model_paths") or []
+    if not paths:
+        log.error("No Bayesian models to evaluate (run training/load_models first).")
+        return
+    _, mc_passes = _eval_cfg(config)
+    _evaluate_ensemble(
+        model_paths=paths,
+        kind=KIND_BAYESIAN,
+        config=config,
+        state=state,
+        log=log,
+        track=track,
+        noise_levels=[0.0],
+        n_passes=mc_passes,
+        out_path=_out_path(config, "bayesian.nc"),
+    )
+
+
+def evaluate_noisy_task(
+    track: ProgressTracker,
+    log: PipelineLogger,
+    config: Dict[str, Any],
+    state: Dict[str, Any],
+) -> None:
+    """Step 6: evaluate both ensembles across the configured Gaussian noise levels."""
+    noise_levels, mc_passes = _eval_cfg(config)
+    normal_paths = state.get("normal_model_paths") or []
+    bayesian_paths = state.get("bayesian_model_paths") or []
+
+    if not normal_paths and not bayesian_paths:
+        log.error("No models to evaluate on noised data.")
+        return
+
+    if normal_paths:
+        _evaluate_ensemble(
+            model_paths=normal_paths,
+            kind=KIND_NORMAL,
+            config=config,
+            state=state,
+            log=log,
+            track=track,
+            noise_levels=noise_levels,
+            n_passes=1,
+            out_path=_out_path(config, "noisy_normal.nc"),
+        )
+    else:
+        log.info("No normal models found; skipping normal noisy evaluation.")
+
+    if bayesian_paths:
+        _evaluate_ensemble(
+            model_paths=bayesian_paths,
+            kind=KIND_BAYESIAN,
+            config=config,
+            state=state,
+            log=log,
+            track=track,
+            noise_levels=noise_levels,
+            n_passes=mc_passes,
+            out_path=_out_path(config, "noisy_bayesian.nc"),
+        )
+    else:
+        log.info("No Bayesian models found; skipping Bayesian noisy evaluation.")

tasks/load_models.py

@@ -0,0 +1,72 @@
+"""Discover saved ensemble checkpoints and record their paths in ``state``.
+
+Training deletes each model from VRAM after saving it, so evaluation cannot reuse
+in-memory models. This task scans ``<work_dir>/{normal,bayesian}_models/`` for
+``model_N.pt`` files and stores ordered path lists; the evaluation tasks then load
+one model at a time. Weights are NOT loaded here (that would defeat the memory
+decoupling) -- only discovered and ordered.
+"""
+
+import pathlib as pl
+import re
+from typing import Any, Dict, List
+
+from model.factory import KIND_BAYESIAN, KIND_NORMAL
+from util.progress import ProgressTracker
+from util.ui_logger import PipelineLogger
+
+# Subdirectory (under work_dir) and state key for each model kind.
+_MODEL_SOURCES = {
+    KIND_NORMAL: ("normal_models", "normal_model_paths"),
+    KIND_BAYESIAN: ("bayesian_models", "bayesian_model_paths"),
+}
+
+_MODEL_FILE_RE = re.compile(r"^model_(\d+)\.pt$")
+
+
+def _discover_models(models_dir: pl.Path) -> List[pl.Path]:
+    """Return ``model_N.pt`` files in ``models_dir`` sorted by N (natural order).
+
+    Only the directory's top level is scanned, so intermediate checkpoints under
+    ``intermediate_models/`` are ignored.
+    """
+    if not models_dir.is_dir():
+        return []
+
+    numbered: List[tuple[int, pl.Path]] = []
+    for entry in models_dir.iterdir():
+        if not entry.is_file():
+            continue
+        match = _MODEL_FILE_RE.match(entry.name)
+        if match:
+            numbered.append((int(match.group(1)), entry))
+
+    numbered.sort(key=lambda pair: pair[0])
+    return [path for _, path in numbered]
+
+
+def load_models_task(
+    track: ProgressTracker,
+    log: PipelineLogger,
+    config: Dict[str, Any],
+    state: Dict[str, Any],
+) -> None:
+    work_dir = pl.Path(config["work_dir"])
+
+    track.update(total=len(_MODEL_SOURCES), advance=0)
+    total_found = 0
+    for kind, (subdir, state_key) in _MODEL_SOURCES.items():
+        paths = _discover_models(work_dir / subdir)
+        state[state_key] = paths
+        total_found += len(paths)
+        if paths:
+            log.info(f"Found {len(paths)} {kind} model(s) in {subdir}/.")
+        else:
+            log.info(f"No {kind} models found in {subdir}/ (evaluation will skip).")
+        track.update(advance=1)
+
+    if total_found == 0:
+        log.error(
+            "No saved models were discovered. Train first, or point the working "
+            "directory at an experiment that contains trained models."
+        )

util/control.py

@@ -0,0 +1,37 @@
+"""Shared pipeline control helpers: cooperative stop/pause and memory release.
+
+These were previously copy-pasted across the training tasks and training loop.
+Centralizing them keeps the stop/pause semantics identical everywhere.
+"""
+
+import gc
+import time
+from threading import Event
+
+import torch
+
+
+def check_control_events(
+    stop_event: Event | None,
+    pause_event: Event | None,
+) -> None:
+    """Cooperatively honor stop/pause events.
+
+    Raises:
+        InterruptedError: if ``stop_event`` is set (including while paused).
+    """
+    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 release_torch_memory(device: str) -> None:
+    """Free Python and accelerator caches between models to limit VRAM growth."""
+    gc.collect()
+    if device.startswith("cuda"):
+        torch.cuda.empty_cache()
+    elif device.startswith("mps") and hasattr(torch, "mps"):
+        torch.mps.empty_cache()