alnn_rewrite/tasks/__init__.py

"""Pipeline task registry and scenario definitions.

A *task* is a callable ``task(tracker, logger, config, state) -> None``. A
*scenario* is an ordered list of task ids the pipeline runs. Tasks communicate
through the mutable ``state`` dict (dataloaders, model paths, control events).

Pipeline stages (see ai/ARCHITECTURE.md):
    1. load_data        -> implemented
    2. train_regular    -> implemented
    3. train_bayesian   -> implemented
    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 . import evaluate
from . import load_data
from . import load_models
from . import train_bayesian
from . import train_normal

PIPELINE_TASKS = {
    "load_data": {
        "task_name": "Load Image and ADNIMERGE",
        "task_func": load_data.load_data_task,
    },
    "train_regular": {
        "task_name": "Train Regular Models",
        "task_func": train_normal.train_normal_task,
    },
    "train_bayesian": {
        "task_name": "Train Bayesian Models",
        "task_func": train_bayesian.train_bayesian_task,
    },
    "load_models": {
        "task_name": "Load Saved Models",
        "task_func": load_models.load_models_task,
    },
    "evaluate_regular": {
        "task_name": "Evaluate Regular Models",
        "task_func": evaluate.evaluate_normal_task,
    },
    "evaluate_bayesian": {
        "task_name": "Evaluate Bayesian Models",
        "task_func": evaluate.evaluate_bayesian_task,
    },
    "evaluate_noisy": {
        "task_name": "Evaluate Models on Noised Data",
        "task_func": evaluate.evaluate_noisy_task,
    },
}

SCENARIOS = {
    "scen_train_all": {
        "label": "1. Train, Evaluate, & Noise Analysis",
        "tasks": [
            "load_data",
            "train_regular",
            "train_bayesian",
            "load_models",
            "evaluate_regular",
            "evaluate_bayesian",
            "evaluate_noisy",
        ],
    },
    "scen_load_all": {
        "label": "2. Load, Evaluate, & Noise Analysis",
        "tasks": [
            "load_data",
            "load_models",
            "evaluate_regular",
            "evaluate_bayesian",
            "evaluate_noisy",
        ],
    },
    "scen_load_eval": {
        "label": "3. Load & Evaluate (Skip Noise)",
        "tasks": [
            "load_data",
            "load_models",
            "evaluate_regular",
            "evaluate_bayesian",
        ],
    },
}