Skip to content
Merged

Dev #85

Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 66 additions & 7 deletions research/vestibular_schwannoma/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ training, inference on new cases, and PACS deployment.
## Contents

- `train_5fold.py`: command-line five-fold training and evaluation.
- `merge_inference_manifests.py`: validate and combine parallel fold subsets for inference.
- `notebooks/01_five_fold_cross_validation.ipynb`: train and compare UNet, DynUNet, and
optional SegMamba models.
- `notebooks/02_inference_new_cases.ipynb`: run one declared model or an explicit ensemble.
Expand Down Expand Up @@ -42,9 +43,62 @@ python train_5fold.py --models unet # One model, all five folds
python train_5fold.py --skip-unavailable
```

The default requests four models across five folds for 500 epochs; run
`python train_5fold.py --help` before starting. For interactive inspection and visualizations,
start Jupyter from this directory or `notebooks/`:
The default requests three models across five folds for 500 epochs; run
`python train_5fold.py --help` before starting. A launcher processes its requested models
and folds sequentially. With five GPUs, run one process per fold, assign each process one
visible GPU, and give it a distinct, previously nonexistent `--results-root`:

```bash
CUDA_VISIBLE_DEVICES=0 python train_5fold.py --models unet --folds 1 --results-root cv_results/unet_fold_1 &
CUDA_VISIBLE_DEVICES=1 python train_5fold.py --models unet --folds 2 --results-root cv_results/unet_fold_2 &
CUDA_VISIBLE_DEVICES=2 python train_5fold.py --models unet --folds 3 --results-root cv_results/unet_fold_3 &
CUDA_VISIBLE_DEVICES=3 python train_5fold.py --models unet --folds 4 --results-root cv_results/unet_fold_4 &
CUDA_VISIBLE_DEVICES=4 python train_5fold.py --models unet --folds 5 --results-root cv_results/unet_fold_5 &
wait
```

Inside each process, its assigned physical GPU is exposed to PyTorch as CUDA device 0.

Before starting parallel jobs, populate `preprocessed/` once with a single process;
concurrent first-time cache creation is not supported. When `MLFLOW_TRACKING_URI` is unset,
processes launched on the same machine from the same fastMONAI checkout automatically share
fastMONAI's repository-root SQLite tracking store
(`sqlite:////absolute/path/to/fastMONAI/mlruns.db`). For multiple machines or a central
tracking service, configure the same remote URI in every shell:

```bash
export MLFLOW_TRACKING_URI=http://mlflow.example:5000
```

Do not merge run IDs from independent local MLflow databases: inference must be able to
resolve every run ID through one tracking URI.

Each launcher atomically updates `completed_run_ids.json` after every successful fold, so
completed work remains mergeable if a later fold is interrupted. A subset job intentionally
does not create `inference_run_ids.json`; its completed registry is also rejected by the
inference loader. Combine disjoint completed-fold registries into a new inference-only
results root. The merger requires exactly folds 1-5, verifies that dataset/splits,
preprocessing, model/loss, and training settings match, and rejects missing or overlapping
folds, duplicate MLflow run IDs, and an existing output root:

```bash
python merge_inference_manifests.py \
cv_results/unet_fold_1/completed_run_ids.json \
cv_results/unet_fold_2/completed_run_ids.json \
cv_results/unet_fold_3/completed_run_ids.json \
cv_results/unet_fold_4/completed_run_ids.json \
cv_results/unet_fold_5/completed_run_ids.json \
--model unet \
--output-root cv_results/unet_5fold_merged
```

To replace only fold 1, train it into a new results root and merge that new
`completed_run_ids.json` with registries containing folds 2-5; omit the old fold-1
registry. The replacement is accepted only when its training contract matches, and the
merged `--output-root` must also be new.

Use `cv_results/unet_5fold_merged/inference_run_ids.json` in notebook 02. For interactive
inspection and visualizations, start Jupyter from this directory or `notebooks/`:

```bash
jupyter lab notebooks/01_five_fold_cross_validation.ipynb
Expand All @@ -65,10 +119,15 @@ to `workflow/`. Training model configs contain the VS-specific architecture and
model reconstruction, patch inference, metrics, and artifact formats remain fastMONAI
responsibilities.

Training retains weights-only `.pth` checkpoints for warm-starting or further fitting with a
newly initialized optimizer and learning-rate schedule; they are not exact training-resume
checkpoints. Inference and deployment use strict-loaded `.safetensors` artifacts. Generated
data, results, tracking stores, checkpoints, and model bundles are excluded from Git.
Training writes selected fold checkpoints below
`<results-root>/<model>/fold_<n>/checkpoints/`, so different folds and models cannot overwrite
each other. All-data learners are independently scoped below
`<results-root>/<model>/all_data/`, but their final artifacts are stored in MLflow rather than as
a local checkpoint. The fold `.pth` files support warm-starting or further fitting with a newly
initialized optimizer and learning-rate schedule; they are not exact training-resume
checkpoints. Final and best inference artifacts remain isolated in their MLflow runs. Inference
and deployment use strict-loaded `.safetensors` artifacts. Generated data, results, tracking
stores, checkpoints, and model bundles are excluded from Git.

For container preparation and execution, see
[deployment/pacs/README.md](deployment/pacs/README.md).
Expand Down
66 changes: 66 additions & 0 deletions research/vestibular_schwannoma/merge_inference_manifests.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
#!/usr/bin/env python3
"""Merge parallel fold-training manifests into one inference selection."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path


PROJECT_ROOT = Path(__file__).resolve().parent
RESEARCH_ROOT = PROJECT_ROOT.parent
if str(RESEARCH_ROOT) not in sys.path:
sys.path.insert(0, str(RESEARCH_ROOT))

from vestibular_schwannoma.workflow.run_selection import ( # noqa: E402
merge_fold_run_selections,
)


def _parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description=(
"Validate and merge disjoint fold inference manifests produced by "
"independently launched training jobs."
)
)
parser.add_argument(
"selection_files",
nargs="+",
type=Path,
help=(
"Source completed_run_ids.json or inference_run_ids.json files to merge."
),
)
parser.add_argument(
"--model",
required=True,
help="Model key whose best fold runs should be merged, for example unet.",
)
parser.add_argument(
"--output-root",
type=Path,
required=True,
help="New directory in which to write the merged inference_run_ids.json.",
)
return parser


def _project_path(path: Path) -> Path:
return path if path.is_absolute() else PROJECT_ROOT / path


def main(argv: list[str] | None = None) -> int:
args = _parser().parse_args(argv)
destination = merge_fold_run_selections(
[_project_path(path) for path in args.selection_files],
model_key=args.model,
output_root=_project_path(args.output_root),
)
print(f"Merged inference run selection: {destination}")
return 0


if __name__ == "__main__":
raise SystemExit(main())
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,9 @@
"\n",
"Patch-loading settings depend on available hardware. More queue workers may improve training throughput, while a larger queue increases RAM usage.\n",
"\n",
"`training_seed` initializes training randomness before each independent run while retaining cuDNN performance optimizations. In an all-data run, every case remains in training and the first case by stable `case_id` order is duplicated only for fastai's validation phase. Its metric is an internal monitor, not held-out evaluation."
"`training_seed` initializes training randomness before each independent run while retaining cuDNN performance optimizations. In an all-data run, every case remains in training and the first case by stable `case_id` order is duplicated only for fastai's validation phase. Its metric is an internal monitor, not held-out evaluation.\n",
"\n",
"Every independently launched job must use a distinct, previously nonexistent `RESULTS_ROOT`. Re-run this configuration cell before starting another sweep. After parallel fold subsets finish (or one is interrupted after completing some folds), use `merge_inference_manifests.py` to validate and combine their `completed_run_ids.json` registries. The merger rejects different dataset, split, preprocessing, model, loss, or training contracts. A partial `completed_run_ids.json` registry cannot be used directly for inference."
]
},
{
Expand All @@ -105,13 +107,13 @@
" use_tta=True,\n",
" compile_models=True,\n",
" target_spacing=(0.4102, 0.4102, 1.5),\n",
" patch_size=(192, 192, 48),\n",
" patch_size=(256, 256, 48),\n",
" queue_num_workers=4,\n",
" queue_length=300,\n",
")\n",
"\n",
"DATA_CSV = \"data/ml_dataset.csv\"\n",
"RESULTS_RUN = datetime.now(timezone.utc).strftime(\"%Y%m%dT%H%M%SZ\")\n",
"RESULTS_RUN = datetime.now(timezone.utc).strftime(\"%Y%m%dT%H%M%S%fZ\")\n",
"RESULTS_ROOT = Path(\"cv_results\") / RESULTS_RUN\n",
"\n",
"torch.backends.cudnn.benchmark = True\n",
Expand Down Expand Up @@ -172,7 +174,9 @@
"source": [
"## 4. Preprocess once to disk\n",
"\n",
"Preprocessing is fold-independent. `preprocess_dataset()` automatically creates a versioned preprocessing cache and a `preprocessing_manifest.json` file. On later runs, fastMONAI uses the manifest to verify that the source files and preprocessing settings are unchanged before reusing the cache. `PatchConfig(preprocessed=True)` prevents preprocessing from being applied twice during training.\n"
"Preprocessing is fold-independent. `preprocess_dataset()` automatically creates a versioned preprocessing cache and a `preprocessing_manifest.json` file. On later runs, fastMONAI uses the manifest to verify that the source files and preprocessing settings are unchanged before reusing the cache. `PatchConfig(preprocessed=True)` prevents preprocessing from being applied twice during training.\n",
"\n",
"Populate a new preprocessing cache with one process before launching parallel training jobs; concurrent first-time cache creation is not supported.\n"
]
},
{
Expand Down Expand Up @@ -543,7 +547,8 @@
"- `cv_results/<RESULTS_RUN>/<model>/fold_N/`: metrics and predictions.\n",
"- `cv_results/<RESULTS_RUN>/<model>/cv_summary.csv`: one complete model summary.\n",
"- `cv_results/<RESULTS_RUN>/cv_model_comparison.csv`: cross-model summary.\n",
"- `cv_results/<RESULTS_RUN>/inference_run_ids.json`: exact completed MLflow runs for notebook 02.\n",
"- `cv_results/<RESULTS_RUN>/completed_run_ids.json`: atomically updated MLflow runs for completed folds, including before a later interruption.\n",
"- `cv_results/<RESULTS_RUN>/inference_run_ids.json`: exact fully completed MLflow runs for notebook 02.\n",
"\n",
"MLflow retains final and selected weights-only checkpoints plus strict-loadable Safetensors artifacts. All-data runs produce only final artifacts.\n",
"\n",
Expand Down
6 changes: 3 additions & 3 deletions research/vestibular_schwannoma/tests/workflow/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ def test_defaults_preserve_the_notebook_experiment(self):
self.assertEqual(config.model_keys, ("unet", "dynunet", "segmamba"))
self.assertEqual(config.folds, (1, 2, 3, 4, 5))
self.assertEqual(config.target_spacing, (0.4102, 0.4102, 1.5))
self.assertEqual(config.patch_size, (192, 192, 48))
self.assertEqual(config.patch_size, (256, 256, 48))
self.assertEqual(config.epochs, 500)
self.assertEqual(config.batch_size, 4)
self.assertEqual(config.training_seed, 42)
Expand Down Expand Up @@ -44,7 +44,7 @@ def test_invalid_declarations_fail_early(self):
{"model_keys": ("unet",), "epochs": 0},
{"model_keys": ("unet",), "training_seed": -1},
{"model_keys": ("unet",), "training_seed": True},
{"model_keys": ("unet",), "patch_size": (192, 192, 0)},
{"model_keys": ("unet",), "patch_size": (256, 256, 0)},
{
"model_keys": ("unet",),
"foreground_sampling_probability": 0,
Expand All @@ -62,7 +62,7 @@ def test_patch_factory_preserves_training_and_inference_contract(self):
config = ExperimentConfig(model_keys=("unet",))
normalization = [ZNormalization(masking_method="foreground")]
patch = make_patch_config(config, normalization)
self.assertEqual(patch.patch_size, [192, 192, 48])
self.assertEqual(patch.patch_size, [256, 256, 48])
self.assertEqual(patch.target_spacing, [0.4102, 0.4102, 1.5])
self.assertEqual(patch.label_probabilities, {0: 0.2, 1: 0.8})
self.assertTrue(patch.preprocessed)
Expand Down
Loading