From 70ef353d8bebb76cf2733fd9570e6b83adbbba77 Mon Sep 17 00:00:00 2001 From: iback Date: Fri, 21 Aug 2026 06:39:23 +0000 Subject: [PATCH 1/2] feat: one trainer, a CLI, and a coverage matrix that cannot go stale The three nnU-Net trainers collapse into one. Which sections a config populates decides whether augmentation runs on the dataloader, on the batch, or both, and pipeline.mode decides how the GPU pipeline is arranged, so the CPU/GPU/Hybrid split no longer needs a class each. The survivor keeps the name nnUNetTrainerDAExtGPU. nnU-Net writes the trainer class name into every checkpoint and resolves the class from it at inference, so renaming it would strand every model already trained. smauglab/cli.py answers, from the registry, the questions that previously needed reading four `if` ladders side by side: smauglab list --backend gpu what a GPU config can name, in pipeline order smauglab matrix which backends implement each augmentation smauglab show RandomScharrGPU one augmentation's parameters and defaults smauglab validate cfg.json every problem in a config, at once smauglab template a config naming every augmentation smauglab hash cfg.json the content-addressed config identity `hash` is the only one that does not need the registry, so it is the only one that does not pay for importing torch. The README coverage matrix and smauglab/configs/all_augmentations.json are generated by `--write` and checked by `--check`, and a test runs `--check`. They are derived from the registry, so they cannot describe augmentations that do not exist or miss ones that do. Two migration rules were missing and are added, with the configs regenerated: * A config carrying a RandomChooseXTransforms block was only ever used with the list trainer, and that choice was the trainer *class*. It becomes pipeline.mode=random_order, or those configs silently switch to the sequential pipeline. * nnU-Net's own SpatialTransform used to be appended by the trainer with its settings hardcoded, so no config named it. The trainer builds its CPU pipeline from the config now, so it has to be in the file or it silently stops running. The block is verbatim from get_training_transforms; patch_size and rotation stay out because nnU-Net supplies them at runtime. Both were caught by the trainer tests rather than by inspection, which is the argument for having them. Co-Authored-By: Claude Opus 5 --- README.md | 44 +++ migration/migrate.py | 25 ++ pyproject.toml | 1 + smauglab/cli.py | 279 ++++++++++++++++ smauglab/configs/transform_params.json | 14 + smauglab/trainers/nnUNetTrainerDAExt.py | 347 +++++++------------- smauglab/trainers/nnUNetTrainerTest.py | 5 + smauglab/trainers/utils.py | 20 +- unit_tests/test_cli.py | 139 ++++++++ unit_tests/test_registered_augmentations.py | 65 ++++ unit_tests/test_trainers.py | 198 +++++++++++ 11 files changed, 888 insertions(+), 249 deletions(-) create mode 100644 smauglab/cli.py create mode 100644 unit_tests/test_cli.py create mode 100644 unit_tests/test_trainers.py diff --git a/README.md b/README.md index e2e1853..ffb35e7 100644 --- a/README.md +++ b/README.md @@ -168,3 +168,47 @@ If you use SmaugLab, please make sure to cite the following paper: year={2026} } ``` + +## Available augmentations + +Which augmentations exist, and which backends implement each one. An empty cell +means no implementation on that backend yet. Regenerate with `smauglab matrix --write`. + + +| Augmentation | Group | GPU | CPU | MONAI | +| --- | --- | --- | --- | --- | +| flip | GEO | `RandomFlipTransformGPU` | — | — | +| affine | GEO | `RandomAffineGPU` | — | — | +| crop | GEO | `RandomCropTransformGPU` | — | — | +| spatial | GEO | — | `SpatialTransform` | — | +| gaussian_noise | GE | `RandomGaussianNoiseGPU` | `GaussianNoiseTransform` | — | +| gaussian_blur | GE | `RandomGaussianBlurGPU` | `GaussianBlurTransform` | — | +| brightness | GE | `RandomBrightnessGPU` | `MultiplicativeBrightnessTransform` | — | +| contrast | GE | `RandomContrastGPU` | `ContrastTransform` | — | +| gamma | GE | `RandomGammaGPU` | `GammaTransform` | — | +| inv_gamma | GE | `RandomInvGammaGPU` | `InvertedGammaTransform` | — | +| clamp | GE | `RandomClampGPU` | — | — | +| low_res | GE | `RandomLowResTransformGPU` | `SimulateLowResolutionTransform` | — | +| acq | GE | `RandomAcqTransformGPU` | — | — | +| zscore | GE | `ZscoreNormalizationGPU` | `ZscoreNormalization` | — | +| mirror | GEO | — | `MirrorTransform` | — | +| scharr | TA | `RandomScharrGPU` | `ScharrConvTransform` | — | +| laplace | TA | `RandomLaplaceGPU` | `LaplaceConvTransform` | — | +| unsharp_mask | TA | `RandomUnsharpMaskGPU` | — | — | +| rand_conv | TA | `RandomRandConvGPU` | — | — | +| bias_field | TA | `RandomBiasFieldGPU` | — | — | +| inverse | TA | `RandomInverseGPU` | — | — | +| histogram_equal | TA | `RandomHistogramEqualizationGPU` | `HistogramEqualTransform` | — | +| redistribute_seg | TA | `RandomRedistributeSegGPU` | `RedistributeTransform` | — | +| palette | TA | `RandomPaletteGPU` | — | — | +| domain_transfer | TA | `RandomDomainTransferGPU` | — | — | +| synthseg | TA | `RandomSynthSegGPU` | — | — | +| artifact | TA | — | `ArtifactTransform` | — | +| spatial_custom | GEO | — | `SpatialCustomTransform` | — | +| shape | GE | — | `ShapeTransform` | — | +| func_log1p | TA | `RandomLog1pGPU` | `Log1pTransform` | — | +| func_sqrt | TA | `RandomSqrtGPU` | `SqrtTransform` | — | +| func_sin | TA | `RandomSinGPU` | `SinTransform` | — | +| func_exp | TA | `RandomExpGPU` | `ExpTransform` | — | +| func_sigmoid | TA | `RandomSigmoidGPU` | `SigmoidTransform` | — | + diff --git a/migration/migrate.py b/migration/migrate.py index bda700a..27b99bf 100644 --- a/migration/migrate.py +++ b/migration/migrate.py @@ -389,6 +389,26 @@ def migrate(payload: dict, source: str = "") -> tuple[dict, list[str]]: target["SpatialTransform"] = _migrate_block("SpatialTransform", spatial, Backend.CPU, source, legacy_key="SpatialTransform") notes.append("moved nnUNetSpatialTransform -> CPU.SpatialTransform") + if backend is Backend.CPU and "SpatialTransform" not in migrated and "SpatialTransform" not in out.get("CPU", {}): + # nnU-Net's own SpatialTransform used to be appended by the trainer with + # these values hardcoded, so no config named it. The trainer builds its CPU + # pipeline from the config now, so it has to be in the file or it silently + # stops running. Values are verbatim from get_training_transforms; + # patch_size and rotation stay out because nnU-Net supplies them at runtime + # (they are context_params on the registry entry). + migrated["SpatialTransform"] = { + "patch_center_dist_from_border": 0, + "random_crop": False, + "p_elastic_deform": 0, + "p_rotation": 0, + "p_scaling": 0, + "scaling": [0.7, 1.4], + "p_synchronize_scaling_across_axes": 1, + "bg_style_seg_sampling": False, + "mode_seg": "nearest", + } + notes.append("added CPU.SpatialTransform, which the trainer used to hardcode") + if backend is Backend.CPU: axes = section.get("mirror_axes") if axes: @@ -398,6 +418,11 @@ def migrate(payload: dict, source: str = "") -> tuple[dict, list[str]]: choose = section.get("RandomChooseXTransforms") if isinstance(choose, dict): out.setdefault("pipeline", {})["random_choose"] = choose + # Which pipeline ran used to be the *trainer class*: a config carrying this + # block was only ever used with the list trainer. That choice is + # pipeline.mode now, so it has to be written down. + out["pipeline"].setdefault("mode", "random_order") + notes.append("moved RandomChooseXTransforms -> pipeline.random_choose, mode=random_order") notes.append("moved RandomChooseXTransforms -> pipeline.random_choose") if migrated: diff --git a/pyproject.toml b/pyproject.toml index a00694a..44d122f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,6 +94,7 @@ all = ["monai", "tqdm", "wandb"] dev = ["build", "coverage", "mypy", "pre-commit", "pytest", "pytest-cov", "ruff", "twine"] [tool.poetry.scripts] +smauglab = "smauglab.cli:main" smauglab_add_nnunettrainer = "smauglab.add_trainer:main" [build-system] diff --git a/smauglab/cli.py b/smauglab/cli.py new file mode 100644 index 0000000..c2cc236 --- /dev/null +++ b/smauglab/cli.py @@ -0,0 +1,279 @@ +"""The `smauglab` command line: what augmentations exist, and are my configs valid? + +Before the registry, both questions could only be answered by reading four +hand-written dispatch ladders side by side. These subcommands read the registry, so +they cannot go out of date: + + smauglab list --backend gpu what a GPU config can name, in pipeline order + smauglab matrix which backends implement each augmentation + smauglab show RandomScharrGPU one augmentation's parameters and defaults + smauglab validate config.json strict check, every problem at once + smauglab template --backend gpu a config naming everything, at defaults + smauglab hash config.json content-addressed config identity + +Bringing a pre-registry config forward is a one-time job and is not a subcommand: the +migrator lives in `migration/` in the repository, not in the wheel. + +`hash` is the only one that does not need the registry; the rest import torch and +kornia to populate it, which takes a few seconds on first use. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +REPO = Path(__file__).resolve().parent.parent +README = REPO / "README.md" +TEMPLATE_PATH = Path(__file__).resolve().parent / "configs" / "all_augmentations.json" +BEGIN = "" +END = "" + + +def _backend(name: str): + from smauglab.registry import Backend + + return Backend[name.upper()] + + +# --- generated artefacts ---------------------------------------------------------- + + +def matrix_block() -> str: + from smauglab import registry + + return f"{BEGIN}\n{registry.render_matrix('md')}\n{END}" + + +def readme_with_matrix(text: str) -> str: + """Replace the marked block, or append the section if it is not there yet.""" + if BEGIN in text and END in text: + head, rest = text.split(BEGIN, 1) + _, tail = rest.split(END, 1) + return f"{head}{matrix_block()}{tail}" + section = ( + "\n## Available augmentations\n\n" + "Which augmentations exist, and which backends implement each one. An empty cell\n" + "means no implementation on that backend yet. Regenerate with `smauglab matrix --write`.\n\n" + matrix_block() + "\n" + ) + return text.rstrip("\n") + "\n" + section + + +def template_json() -> str: + from smauglab import registry + from smauglab.registry import Backend + + payload = {b.value: registry.render_template(b) for b in (Backend.GPU, Backend.CPU)} + return json.dumps(payload, indent=4) + "\n" + + +def _sync(targets: dict[Path, str], *, write: bool, check: bool) -> int: + stale = [path for path, content in targets.items() if not path.is_file() or path.read_text() != content] + if check: + for path in stale: + print(f"out of date: {path.relative_to(REPO)}") + if stale: + print("Run `smauglab matrix --write` / `smauglab template --write` and commit the result.") + return 1 if stale else 0 + if write: + for path in stale: + path.write_text(targets[path]) + print(f"wrote {path.relative_to(REPO)}") + if not stale: + print("already up to date") + return 0 + + +# --- subcommands ------------------------------------------------------------------ + + +def cmd_list(args) -> int: + from smauglab import registry + from smauglab.registry import AugType + + backend = _backend(args.backend) if args.backend else None + group = AugType[args.group.upper()] if args.group else None + entries = registry.entries(backend=backend, group=group) + + if args.json: + print( + json.dumps( + [ + { + "name": e.name, + "backend": e.backend.value, + "aug_id": e.aug_id.value, + "group": e.group.value, + "position": registry.pipeline_position(e), + "summary": e.summary, + } + for e in entries + ], + indent=2, + ) + ) + return 0 + + if not entries: + print("no augmentations match") + return 0 + width = max(len(e.name) for e in entries) + current = None + for entry in entries: + if entry.backend is not current: + current = entry.backend + print(f"\n{current.value} ({len([e for e in entries if e.backend is current])}), in pipeline order:") + print(f" {registry.pipeline_position(entry):>4} {entry.name:<{width}} {entry.group.value:<4} {entry.summary}") + return 0 + + +def cmd_matrix(args) -> int: + from smauglab import registry + + if args.write or args.check: + return _sync( + {README: readme_with_matrix(README.read_text()), TEMPLATE_PATH: template_json()}, + write=args.write, + check=args.check, + ) + if args.format == "json": + table = registry.matrix() + print( + json.dumps( + {aug_id.value: {b.value: (e.name if e else None) for b, e in row.items()} for aug_id, row in table.items()}, + indent=2, + ) + ) + return 0 + print(registry.render_matrix(args.format)) + return 0 + + +def cmd_show(args) -> int: + from smauglab import registry + + try: + entry = registry.get(args.name) + except registry.UnknownAugmentationError as exc: + print(str(exc), file=sys.stderr) + return 1 + + print( + f"{entry.name} ({entry.backend.value}, group {entry.group.value}, position {registry.pipeline_position(entry)}, aug_id {entry.aug_id.value})" + ) + if entry.summary: + print(f" {entry.summary}") + if entry.forwards_to is not None: + print(f" forwards extra parameters to {entry.forwards_to.__name__}") + if entry.external_asset: + print(f" needs an external asset; set ${entry.external_asset}") + print(f"\n module: {entry.cls.__module__}") + + params = registry.accepted_params(entry) + required = registry.required_params(entry) + print(f"\n parameters ({len(params)}):") + for name in sorted(params): + default = params[name].default + shown = "REQUIRED" if name in required else repr(default) + print(f" {name:<34} {shown}") + if entry.context_params: + print(f"\n supplied by the trainer, not the config: {', '.join(entry.context_params)}") + return 0 + + +def cmd_validate(args) -> int: + from smauglab.config import validate_file + + failed = 0 + for path in args.configs: + problems = validate_file(path) + if problems: + failed += 1 + print(f"{path}: {len(problems)} problem(s)") + for problem in problems: + print(f" - {problem}") + elif not args.quiet: + print(f"{path}: ok") + return 1 if failed else 0 + + +def cmd_template(args) -> int: + from smauglab import registry + + if args.write or args.check: + return _sync({TEMPLATE_PATH: template_json()}, write=args.write, check=args.check) + section = registry.render_template(_backend(args.backend)) + text = json.dumps({args.backend.upper(): section}, indent=4) + "\n" + if args.output: + Path(args.output).write_text(text) + print(f"wrote {args.output}") + else: + print(text, end="") + return 0 + + +def cmd_hash(args) -> int: + from smauglab.config import config_hash + + for path in args.configs: + payload = json.loads(Path(path).read_text()) + print(f"{config_hash(payload)[:8]} {path}") + return 0 + + +# --- wiring ----------------------------------------------------------------------- + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="smauglab", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = parser.add_subparsers(dest="command", required=True) + + p = sub.add_parser("list", help="registered augmentations, in pipeline order") + p.add_argument("--backend", choices=["gpu", "cpu", "monai"]) + p.add_argument("--group", choices=["geo", "ge", "ta"]) + p.add_argument("--json", action="store_true") + p.set_defaults(func=cmd_list) + + p = sub.add_parser("matrix", help="which backends implement each augmentation") + p.add_argument("--format", choices=["md", "table", "json"], default="table") + p.add_argument("--write", action="store_true", help="update README.md and the template config") + p.add_argument("--check", action="store_true", help="exit 1 if either is out of date") + p.set_defaults(func=cmd_matrix) + + p = sub.add_parser("show", help="one augmentation's parameters and defaults") + p.add_argument("name") + p.set_defaults(func=cmd_show) + + p = sub.add_parser("validate", help="strict config check, reporting every problem") + p.add_argument("configs", nargs="+") + p.add_argument("-q", "--quiet", action="store_true", help="only report failures") + p.set_defaults(func=cmd_validate) + + p = sub.add_parser("template", help="a config naming every augmentation at its defaults") + p.add_argument("--backend", choices=["gpu", "cpu"], default="gpu") + p.add_argument("-o", "--output") + p.add_argument("--write", action="store_true", help="update the shipped template config") + p.add_argument("--check", action="store_true", help="exit 1 if it is out of date") + p.set_defaults(func=cmd_template) + + p = sub.add_parser("hash", help="content-addressed config identity") + p.add_argument("configs", nargs="+") + p.set_defaults(func=cmd_hash) + + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + if getattr(args, "output", None) and len(getattr(args, "configs", [])) > 1: + print("-o takes a single input config", file=sys.stderr) + return 2 + result: Any = args.func(args) + return int(result or 0) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/smauglab/configs/transform_params.json b/smauglab/configs/transform_params.json index 621c5ff..bbaaeaf 100644 --- a/smauglab/configs/transform_params.json +++ b/smauglab/configs/transform_params.json @@ -133,6 +133,20 @@ "p_retain_stats": 1, "p": 0.3 }, + "SpatialTransform": { + "patch_center_dist_from_border": 0, + "random_crop": false, + "p_elastic_deform": 0, + "p_rotation": 0, + "p_scaling": 0, + "scaling": [ + 0.7, + 1.4 + ], + "p_synchronize_scaling_across_axes": 1, + "bg_style_seg_sampling": false, + "mode_seg": "nearest" + }, "MirrorTransform": { "allowed_axes": [ 0, diff --git a/smauglab/trainers/nnUNetTrainerDAExt.py b/smauglab/trainers/nnUNetTrainerDAExt.py index 47098ab..f1ae88e 100644 --- a/smauglab/trainers/nnUNetTrainerDAExt.py +++ b/smauglab/trainers/nnUNetTrainerDAExt.py @@ -1,7 +1,25 @@ +"""The SmaugLab nnU-Net trainer. + +One class, because the config already says everything the three previous trainers +encoded between them. `nnUNetTrainerDAExt` built the CPU pipeline, `...GPU` built the +GPU one plus nnU-Net's SpatialTransform, and `...Hybrid` built both -- but a config is +sectioned into "CPU" and "GPU", so which sections are populated decides that on its +own: + + transform_params.json CPU: 19 GPU: 0 -> CPU-only, as ...DAExt did + transform_params_gpu.json CPU: 1 GPU: 26 -> GPU-only, as ...DAExtGPU did + transform_params_hybrid.json CPU: 19 GPU: 24 -> both, as ...DAExtHybrid did + +The class keeps the name `nnUNetTrainerDAExtGPU` whatever the config contains. That is +not cosmetic: nnU-Net writes the trainer class name into every checkpoint +(`checkpoint['trainer_name']`) and resolves the class from it at inference, so +renaming it would make several hundred trained models unloadable. +""" + import importlib -import json import os import shutil +import warnings from typing import Union import numpy as np @@ -9,10 +27,8 @@ from batchgeneratorsv2.helpers.scalar_type import RandomScalar from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform from batchgeneratorsv2.transforms.nnunet.seg_to_onehot import MoveSegAsOneHotToDataTransform -from batchgeneratorsv2.transforms.spatial.spatial import SpatialTransform from batchgeneratorsv2.transforms.utils.compose import ComposeTransforms from batchgeneratorsv2.transforms.utils.deep_supervision_downsampling import DownsampleSegForDSTransform -from batchgeneratorsv2.transforms.utils.pseudo2d import Convert2DTo3DTransform, Convert3DTo2DTransform from batchgeneratorsv2.transforms.utils.remove_label import RemoveLabelTansform from batchgeneratorsv2.transforms.utils.seg_to_regions import ConvertSegmentationToRegionsTransform from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer @@ -20,114 +36,79 @@ from torch import autocast from smauglab import configs +from smauglab.config import load_config +from smauglab.registry import Backend from smauglab.trainers.utils import DownsampleSegForDSTransformCustom, nnunet_tail_transforms -from smauglab.transforms.cpu.transforms import AugTransforms -from smauglab.transforms.gpu.transforms import AugTransformsGPU - - -class nnUNetTrainerDAExt(nnUNetTrainer): - def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): - super().__init__(plans, configuration, fold, dataset_json, device) - - def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): - rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = ( - super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() +from smauglab.transforms.build import build_cpu_pipeline, build_gpu_pipeline +from smauglab.transforms.gpu.base import AugmentationSequentialCustom + +#: Env var naming the config to train with. +CONFIG_ENV = "SMAUGLAB_PARAMS_JSON" + +#: Previous name, still honoured so existing sweep scripts keep working. The three +#: old variables (_CPU_JSON, _GPU_JSON, _HYBRID_JSON) picked a trainer as much as a +#: file; only this one ever had an external caller. +LEGACY_CONFIG_ENV = "SMAUGLAB_PARAMS_GPU_JSON" + +DEFAULT_CONFIG = "transform_params_gpu.json" + + +def resolve_config_path() -> str: + """Locate the config: new env var, then the deprecated one, then the default.""" + path = os.environ.get(CONFIG_ENV) + if path: + return path + legacy = os.environ.get(LEGACY_CONFIG_ENV) + if legacy: + warnings.warn( + f"{LEGACY_CONFIG_ENV} is deprecated; use {CONFIG_ENV}. The config's CPU and GPU " + "sections now decide which augmentations run, so the name no longer implies a backend.", + DeprecationWarning, + stacklevel=2, ) - # Remove mirroring - mirror_axes = None - self.inference_allowed_mirroring_axes = None - return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes + return legacy + return str(importlib.resources.files(configs) / DEFAULT_CONFIG) - @staticmethod - def get_training_transforms( - patch_size: Union[np.ndarray, tuple[int, ...]], - rotation_for_DA: RandomScalar, - deep_supervision_scales: Union[list, tuple, None], - mirror_axes: tuple[int, ...], - do_dummy_2d_data_aug: bool, - use_mask_for_norm: list[bool] | None = None, - is_cascaded: bool = False, - foreground_labels: Union[tuple[int, ...], list[int]] | None = None, - regions: list[Union[list[int], tuple[int, ...], int]] | None = None, - ignore_label: int | None = None, - retain_stats: bool = False, - ) -> BasicTransform: - transforms = [] - ### Adds transforms - # Load transform parameters from json file - configs_path = importlib.resources.files(configs) - json_path = os.environ.get("SMAUGLAB_PARAMS_CPU_JSON", str(configs_path / "transform_params.json")) - transforms.append( - AugTransforms( - json_path=json_path, - do_dummy_2d_data_aug=do_dummy_2d_data_aug, - patch_size=patch_size, - rotation_for_DA=rotation_for_DA, - mirror_axes=mirror_axes, - ) - ) +def _has_gpu_augmentations(config) -> bool: + """Whether this config asks for anything on the GPU side.""" + return bool(config.names(Backend.GPU)) - if do_dummy_2d_data_aug: - transforms.append(Convert3DTo2DTransform()) - patch_size_spatial = patch_size[1:] - else: - patch_size_spatial = patch_size - transforms.append( - SpatialTransform( - patch_size_spatial, - patch_center_dist_from_border=0, - random_crop=False, - p_elastic_deform=0, - p_rotation=0, - rotation=rotation_for_DA, - p_scaling=0, - scaling=(0.7, 1.4), - p_synchronize_scaling_across_axes=1, - bg_style_seg_sampling=False, - mode_seg="nearest", - ) - ) - - if do_dummy_2d_data_aug: - transforms.append(Convert2DTo3DTransform()) - transforms.extend( - nnunet_tail_transforms( - use_mask_for_norm=use_mask_for_norm, - deep_supervision_scales=deep_supervision_scales, - is_cascaded=is_cascaded, - foreground_labels=foreground_labels, - regions=regions, - ignore_label=ignore_label, - ) - ) - return ComposeTransforms(transforms) +class nnUNetTrainerDAExtGPU(nnUNetTrainer): + """nnU-Net trainer driven entirely by a SmaugLab config. + CPU-section augmentations run in the dataloader worker; GPU-section ones run on + the batch in `train_step`. + """ -class nnUNetTrainerDAExtGPU(nnUNetTrainer): def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): super().__init__(plans, configuration, fold, dataset_json, device) - self.num_epochs = 1000 - - # Load transform parameters from json file - configs_path = importlib.resources.files(configs) - json_path = os.environ.get("SMAUGLAB_PARAMS_GPU_JSON", str(configs_path / "transform_params_gpu.json")) - self.transforms = AugTransformsGPU(json_path=json_path).to(self.device) - print(f"Using SmaugLab GPU transforms with parameters from: {json_path}") - - # Copy json transfrom parameters to output folder - shutil.copy(json_path, os.path.join(self.output_folder, "transform_params_gpu_used_for_training.json")) - - def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): - rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = ( - super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() + json_path = resolve_config_path() + config = load_config(json_path) + + # Built only when the config actually asks for GPU augmentations, so a + # CPU-only config costs nothing per step and train_step stays a no-op. + self.transforms: AugmentationSequentialCustom | None = None + if _has_gpu_augmentations(config): + self.transforms = AugmentationSequentialCustom( + *build_gpu_pipeline( + config.section(Backend.GPU), + mode=config.pipeline_mode(), + options=config.pipeline_options("random_choose"), + source=config.source, + ), + data_keys=["input", "mask"], + same_on_batch=True, + ).to(self.device) + + print(f"Using SmaugLab transforms from: {json_path}") + print( + f" CPU: {len(config.names(Backend.CPU))} augmentations, GPU: {len(config.names(Backend.GPU))}, mode: {config.pipeline_mode().value}" ) - # Remove mirroring - mirror_axes = None - self.inference_allowed_mirroring_axes = None - return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes + + shutil.copy(json_path, os.path.join(self.output_folder, "transform_params_used_for_training.json")) @staticmethod def get_training_transforms( @@ -141,53 +122,41 @@ def get_training_transforms( foreground_labels: Union[tuple[int, ...], list[int]] | None = None, regions: list[Union[list[int], tuple[int, ...], int]] | None = None, ignore_label: int | None = None, - retain_stats: bool = False, ) -> BasicTransform: - transforms = [] - - configs_path = importlib.resources.files(configs) - json_path = os.environ.get("SMAUGLAB_PARAMS_GPU_JSON", str(configs_path / "transform_params_gpu.json")) - with open(json_path) as f: - config = json.load(f) + """Dataloader-side augmentations: whatever the config's CPU section names. - ### Keep some nnunet transforms - if do_dummy_2d_data_aug: - transforms.append(Convert3DTo2DTransform()) - patch_size_spatial = patch_size[1:] - else: - patch_size_spatial = patch_size - - spatial_params = config.get("nnUNetSpatialTransform", {}) + A staticmethod because that is nnU-Net's contract, so it cannot reach the + instance's parsed config and resolves the path itself. `load_config` is + cached, so the file is still read and validated once. + """ + transforms = [] - transforms.append( - SpatialTransform( - patch_size_spatial, - patch_center_dist_from_border=spatial_params.get("patch_center_dist_from_border", 0), - random_crop=spatial_params.get("random_crop", False), - p_elastic_deform=spatial_params.get("p_elastic_deform", 0), - p_rotation=spatial_params.get("p_rotation", 0), + config = load_config(resolve_config_path()) + transforms.extend( + build_cpu_pipeline( + config.section(Backend.CPU), + do_dummy_2d_data_aug=do_dummy_2d_data_aug, + patch_size=patch_size, rotation=rotation_for_DA, - p_scaling=spatial_params.get("p_scaling", 0), - scaling=spatial_params.get("scaling", (0.7, 1.4)), - p_synchronize_scaling_across_axes=spatial_params.get("p_synchronize_scaling_across_axes", 1), - bg_style_seg_sampling=False, - mode_seg="nearest", + source=config.source, ) ) - if do_dummy_2d_data_aug: - transforms.append(Convert2DTo3DTransform()) - + # Deep supervision has to come after whatever last deformed the mask. With GPU + # augmentations that is train_step, so the downsampling happens there; without + # them nothing touches the mask after this point and it belongs here, which is + # where nnU-Net puts it. Passing None is how the tail is told to skip it. transforms.extend( nnunet_tail_transforms( use_mask_for_norm=use_mask_for_norm, - deep_supervision_scales=None, + deep_supervision_scales=None if _has_gpu_augmentations(config) else deep_supervision_scales, is_cascaded=is_cascaded, foreground_labels=foreground_labels, regions=regions, ignore_label=ignore_label, ) ) + return ComposeTransforms(transforms) @staticmethod @@ -220,63 +189,6 @@ def get_validation_transforms( transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales)) return ComposeTransforms(transforms) - def train_step(self, batch: dict) -> dict: - data = batch["data"] - target = batch["target"] - - data = data.to(self.device, non_blocking=True) - # Now target should be a single tensor, not a list - target = target.to(self.device, non_blocking=True) - # if isinstance(target, list): - # target = [i.to(self.device, non_blocking=True) for i in target] - # else: - # target = target.to(self.device, non_blocking=True) - - self.optimizer.zero_grad(set_to_none=True) - # Autocast can be annoying - # If the device_type is 'cpu' then it's slow as heck and needs to be disabled. - # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False) - # So autocast will only be active if we have a cuda device. - with autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context(): - # Apply GPU augmentations to full-resolution data/target - data, target = self.transforms(data, target) - - # Create multi-scale targets for deep supervision after augmentation - deep_supervision_scales = self._get_deep_supervision_scales() - if deep_supervision_scales is not None: - ds_transform = DownsampleSegForDSTransformCustom(ds_scales=deep_supervision_scales) - target = ds_transform(target) - - output = self.network(data) - # del data - l = self.loss(output, target) - - if self.grad_scaler is not None: - self.grad_scaler.scale(l).backward() - self.grad_scaler.unscale_(self.optimizer) - torch.nn.utils.clip_grad_norm_(self.network.parameters(), 12) - self.grad_scaler.step(self.optimizer) - self.grad_scaler.update() - else: - l.backward() - torch.nn.utils.clip_grad_norm_(self.network.parameters(), 12) - self.optimizer.step() - return {"loss": l.detach().cpu().numpy()} - - -class nnUNetTrainerDAExtHybrid(nnUNetTrainer): - def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): - super().__init__(plans, configuration, fold, dataset_json, device) - - # Load transform parameters from json file - configs_path = importlib.resources.files(configs) - json_path = os.environ.get("SMAUGLAB_PARAMS_HYBRID_JSON", str(configs_path / "transform_params_hybrid.json")) - self.transforms = AugTransformsGPU(json_path=json_path).to(self.device) - print(f"Using SmaugLab hybrid transforms with parameters from: {json_path}") - - # Copy json transfrom parameters to output folder - shutil.copy(json_path, os.path.join(self.output_folder, "transform_params_hybrid_used_for_training.json")) - def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = ( super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() @@ -286,48 +198,6 @@ def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): self.inference_allowed_mirroring_axes = None return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes - @staticmethod - def get_training_transforms( - patch_size: Union[np.ndarray, tuple[int, ...]], - rotation_for_DA: RandomScalar, - deep_supervision_scales: Union[list, tuple, None], - mirror_axes: tuple[int, ...], - do_dummy_2d_data_aug: bool, - use_mask_for_norm: list[bool] | None = None, - is_cascaded: bool = False, - foreground_labels: Union[tuple[int, ...], list[int]] | None = None, - regions: list[Union[list[int], tuple[int, ...], int]] | None = None, - ignore_label: int | None = None, - retain_stats: bool = False, - ) -> BasicTransform: - transforms = [] - - ### Adds transforms - # Load transform parameters from json file - configs_path = importlib.resources.files(configs) - json_path = os.environ.get("SMAUGLAB_PARAMS_HYBRID_JSON", str(configs_path / "transform_params_hybrid.json")) - transforms.append( - AugTransforms( - json_path=json_path, - do_dummy_2d_data_aug=do_dummy_2d_data_aug, - patch_size=patch_size, - rotation_for_DA=rotation_for_DA, - mirror_axes=mirror_axes, - ) - ) - - transforms.extend( - nnunet_tail_transforms( - use_mask_for_norm=use_mask_for_norm, - deep_supervision_scales=None, - is_cascaded=is_cascaded, - foreground_labels=foreground_labels, - regions=regions, - ignore_label=ignore_label, - ) - ) - return ComposeTransforms(transforms) - def train_step(self, batch: dict) -> dict: data = batch["data"] target = batch["target"] @@ -346,14 +216,17 @@ def train_step(self, batch: dict) -> dict: # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False) # So autocast will only be active if we have a cuda device. with autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context(): - # Apply GPU augmentations to full-resolution data/target - data, target = self.transforms(data, target) - - # Create multi-scale targets for deep supervision after augmentation - deep_supervision_scales = self._get_deep_supervision_scales() - if deep_supervision_scales is not None: - ds_transform = DownsampleSegForDSTransformCustom(ds_scales=deep_supervision_scales) - target = ds_transform(target) + # Apply GPU augmentations to full-resolution data/target, then build the + # deep-supervision targets from the *augmented* mask. A CPU-only config + # builds no GPU pipeline; nothing has touched the mask since the + # dataloader, which already produced those targets. + if self.transforms is not None: + data, target = self.transforms(data, target) + + deep_supervision_scales = self._get_deep_supervision_scales() + if deep_supervision_scales is not None: + ds_transform = DownsampleSegForDSTransformCustom(ds_scales=deep_supervision_scales) + target = ds_transform(target) output = self.network(data) # del data diff --git a/smauglab/trainers/nnUNetTrainerTest.py b/smauglab/trainers/nnUNetTrainerTest.py index a51f32f..f1fbebd 100644 --- a/smauglab/trainers/nnUNetTrainerTest.py +++ b/smauglab/trainers/nnUNetTrainerTest.py @@ -76,6 +76,7 @@ def get_training_transforms( ignore_label=ignore_label, ) ) + return ComposeTransforms(transforms) @@ -129,6 +130,9 @@ def get_training_transforms( if do_dummy_2d_data_aug: transforms.append(Convert2DTo3DTransform()) + # deep_supervision_scales=None: this trainer always runs GPU augmentations, so + # the mask is still being deformed after this point and train_step builds the + # multi-scale targets from the augmented mask instead. transforms.extend( nnunet_tail_transforms( use_mask_for_norm=use_mask_for_norm, @@ -139,6 +143,7 @@ def get_training_transforms( ignore_label=ignore_label, ) ) + return ComposeTransforms(transforms) def train_step(self, batch: dict) -> dict: diff --git a/smauglab/trainers/utils.py b/smauglab/trainers/utils.py index e7a511e..96d952f 100644 --- a/smauglab/trainers/utils.py +++ b/smauglab/trainers/utils.py @@ -16,18 +16,14 @@ def nnunet_tail_transforms( ) -> list[Any]: """The nnU-Net transforms that follow SmaugLab's augmentations, in order. - Five `get_training_transforms` methods across two modules ended with a - character-identical copy of this -- intensity masking, the -1 label removal, the - two cascade transforms, region conversion and deep-supervision downsampling. - - `deep_supervision_scales=None` skips the downsampling, which is how the GPU - trainers had it: they carry the block commented out, because with GPU - augmentations the mask is still being deformed after this point and the - multi-scale targets have to be built from the augmented mask in `train_step`. - - `get_validation_transforms` deliberately does not use this. Its cascade branch - adds only MoveSegAsOneHotToDataTransform, without the two RandomTransform - wrappers, so it is a different sequence rather than another copy of this one. + Every `get_training_transforms` in this package ended with a character-identical + copy of this -- three of them -- covering intensity masking, the -1 label removal, + the two cascade transforms, region conversion and deep-supervision downsampling. + + `deep_supervision_scales=None` skips the downsampling, which is how + `nnUNetTrainerDAExtGPU` defers it to `train_step`: with GPU augmentations the mask + is still being deformed after this point, so the multi-scale targets have to be + built from the augmented mask instead. """ from batchgeneratorsv2.transforms.nnunet.random_binary_operator import ApplyRandomBinaryOperatorTransform from batchgeneratorsv2.transforms.nnunet.remove_connected_components import ( diff --git a/unit_tests/test_cli.py b/unit_tests/test_cli.py new file mode 100644 index 0000000..a3f134c --- /dev/null +++ b/unit_tests/test_cli.py @@ -0,0 +1,139 @@ +"""The `smauglab` command line. + +Driven through `cli.main` with captured stdout rather than as a subprocess: the +registry import costs a few seconds, and paying it once per process keeps the suite +fast enough to gate every pull request. +""" + +from __future__ import annotations + +import contextlib +import io +import json +import tempfile +import unittest +from pathlib import Path + +from smauglab import cli + +REPO = Path(__file__).resolve().parent.parent +DEFAULT_GPU = REPO / "smauglab" / "configs" / "transform_params_gpu_default01-23.json" +LEGACY = REPO / "unit_tests" / "fixtures" / "legacy_configs" / "configs" / "transform_params_gpu_default01-23.json" + + +def run(*argv: str) -> tuple[int, str]: + out = io.StringIO() + with contextlib.redirect_stdout(out), contextlib.redirect_stderr(out): + code = cli.main(list(argv)) + return code, out.getvalue() + + +class TestList(unittest.TestCase): + def test_lists_gpu_augmentations_in_pipeline_order(self): + code, out = run("list", "--backend", "gpu") + self.assertEqual(code, 0) + self.assertIn("RandomFlipTransformGPU", out) + # order column ascends + orders = [int(line.split()[0]) for line in out.splitlines() if line.startswith(" ") and line.split()[0].isdigit()] + self.assertEqual(orders, sorted(orders)) + + def test_group_filter(self): + _, out = run("list", "--backend", "gpu", "--group", "geo") + self.assertIn("RandomFlipTransformGPU", out) + self.assertNotIn("RandomScharrGPU", out) + + def test_json_output_is_machine_readable(self): + _, out = run("list", "--backend", "cpu", "--json") + payload = json.loads(out) + self.assertTrue(all({"name", "backend", "aug_id", "group", "position"} <= set(e) for e in payload)) + + +class TestMatrix(unittest.TestCase): + def test_markdown_shows_the_monai_column_as_empty(self): + code, out = run("matrix", "--format", "md") + self.assertEqual(code, 0) + self.assertIn("| Augmentation | Group | GPU | CPU | MONAI |", out) + + def test_json_form_reports_missing_backends_as_null(self): + _, out = run("matrix", "--format", "json") + payload = json.loads(out) + self.assertIsNone(payload["palette"]["CPU"]) + self.assertEqual(payload["palette"]["GPU"], "RandomPaletteGPU") + self.assertTrue(all(row["MONAI"] is None for row in payload.values())) + + def test_check_passes_on_a_clean_tree(self): + code, _ = run("matrix", "--check") + self.assertEqual(code, 0) + + +class TestShow(unittest.TestCase): + def test_reports_parameters_and_defaults(self): + code, out = run("show", "RandomScharrGPU") + self.assertEqual(code, 0) + self.assertIn("group TA", out) + self.assertIn("absolute", out) + + def test_marks_required_parameters(self): + _, out = run("show", "MirrorTransform") + self.assertIn("REQUIRED", out) + + def test_flags_trainer_supplied_parameters(self): + _, out = run("show", "SpatialTransform") + self.assertIn("supplied by the trainer", out) + + def test_unknown_name_exits_nonzero_with_a_suggestion(self): + code, out = run("show", "ScharrTransform") + self.assertEqual(code, 1) + self.assertIn("RandomScharrGPU", out) + + +class TestValidate(unittest.TestCase): + def test_shipped_config_is_valid(self): + code, out = run("validate", str(DEFAULT_GPU)) + self.assertEqual(code, 0) + self.assertIn("ok", out) + + def test_a_broken_config_exits_nonzero_and_reports_every_problem(self): + payload = {"GPU": {"ScharrTransform": {"p": 0.1}, "RandomGaussianNoiseGPU": {"probability": 0.2}}} + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "broken.json" + path.write_text(json.dumps(payload)) + code, out = run("validate", str(path)) + self.assertEqual(code, 1) + self.assertIn("2 problem(s)", out) + self.assertIn("RandomScharrGPU", out) + self.assertIn("'probability' -> p", out) + + def test_a_legacy_config_is_rejected_and_points_at_migrate(self): + code, out = run("validate", str(LEGACY)) + self.assertEqual(code, 1) + self.assertIn("migration/", out) + + +class TestTemplateAndHash(unittest.TestCase): + def test_template_names_every_gpu_augmentation(self): + from smauglab import registry + from smauglab.registry import Backend + + _, out = run("template", "--backend", "gpu") + self.assertEqual(set(json.loads(out)["GPU"]), set(registry.names(Backend.GPU))) + + def test_template_round_trips_through_validation(self): + """A template that cannot be loaded would be worse than no template.""" + from smauglab.config import validate_file + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "template.json" + run("template", "--backend", "gpu", "-o", str(path)) + self.assertEqual(validate_file(path), []) + + def test_hash_is_stable_and_content_addressed(self): + code, first = run("hash", str(DEFAULT_GPU)) + self.assertEqual(code, 0) + _, second = run("hash", str(DEFAULT_GPU)) + self.assertEqual(first, second) + self.assertEqual(len(first.split()[0]), 8) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_registered_augmentations.py b/unit_tests/test_registered_augmentations.py index 0a4bf9f..f2774cb 100644 --- a/unit_tests/test_registered_augmentations.py +++ b/unit_tests/test_registered_augmentations.py @@ -11,6 +11,10 @@ from __future__ import annotations +import inspect +import json +import subprocess +import sys import unittest from pathlib import Path @@ -113,3 +117,64 @@ def test_constructible_with_declared_defaults(self): ) continue entry.cls(**dict(entry.smoke_kwargs)) + + +class TestGeneratedArtifactsAreCurrent(unittest.TestCase): + """`smauglab matrix --check` is the anti-staleness guarantee. + + Run as a subprocess rather than re-deriving the comparison here, so the test + exercises the same code path a developer and CI run. + """ + + def test_readme_matrix_and_template_are_up_to_date(self): + result = subprocess.run( + [sys.executable, "-m", "smauglab.cli", "matrix", "--check"], + capture_output=True, + text=True, + cwd=REPO, + check=False, # a non-zero exit is the assertion below, not an error here + ) + self.assertEqual( + result.returncode, + 0, + f"generated artifacts are stale:\n{result.stdout}{result.stderr}", + ) + + def test_template_covers_exactly_the_registered_augmentations(self): + payload = json.loads(TEMPLATE.read_text()) + for backend in (Backend.GPU, Backend.CPU): + with self.subTest(backend=backend.value): + self.assertEqual(set(payload[backend.value]), set(registry.names(backend))) + + def test_template_parameters_are_all_accepted(self): + """Every key in the template must survive the validation stage 6 will apply.""" + payload = json.loads(TEMPLATE.read_text()) + for backend in (Backend.GPU, Backend.CPU): + for name, params in payload[backend.value].items(): + entry = registry.get(name, backend) + accepted = set(registry.accepted_params(entry)) + with self.subTest(entry=f"{backend.value}.{name}"): + self.assertEqual(set(params) - accepted, set()) + + def test_forced_parameters_are_not_offered(self): + """RandomSynthSegGPU overrides these internally; offering them would lie.""" + payload = json.loads(TEMPLATE.read_text()) + synthseg = payload["GPU"]["RandomSynthSegGPU"] + for forced in ("apply_affine", "apply_nonlinear", "flipping", "output_shape"): + with self.subTest(param=forced): + self.assertNotIn(forced, synthseg) + + +class TestForwardedSignatures(unittest.TestCase): + def test_synthseg_accepts_its_generator_parameters(self): + """forwards_to unions the two signatures, replacing a hand-listed allowlist.""" + from smauglab.transforms.synthseg.generator import SynthSegGenerator + + entry = registry.get("RandomSynthSegGPU", Backend.GPU) + accepted = set(registry.accepted_params(entry)) + generator = set(inspect.signature(SynthSegGenerator).parameters) - set(entry.context_params) + self.assertEqual(generator - accepted, set(), "generator parameters missing from the accepted set") + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_trainers.py b/unit_tests/test_trainers.py new file mode 100644 index 0000000..7f67e5d --- /dev/null +++ b/unit_tests/test_trainers.py @@ -0,0 +1,198 @@ +"""The nnU-Net trainer, which is now one class driven entirely by the config. + +Three trainers used to encode the CPU/GPU split in their class names. The config +already carries it -- which sections are populated says what runs -- so there is one +class, and these tests pin that it reproduces what each of the three used to build. + +`get_training_transforms` is a staticmethod (nnU-Net's contract), so it can be driven +directly: no plans.json, no dataset, no GPU. +""" + +from __future__ import annotations + +import importlib.util +import os +import unittest +import warnings +from pathlib import Path + +from batchgeneratorsv2.transforms.utils.compose import ComposeTransforms + +from smauglab.config import load_config +from smauglab.registry import Backend +from smauglab.transforms.build import PipelineMode, build_cpu_pipeline + +REPO = Path(__file__).resolve().parent.parent +CONFIGS = REPO / "smauglab" / "configs" +PATCH = (24, 24, 24) +ROTATION = (-10, 10) +DS_SCALES = [[1, 1, 1], [0.5, 0.5, 0.5]] + +# Raised at import so it works under pytest and `python -m unittest` alike; a +# `pytestmark` would only be understood by one of them. +if importlib.util.find_spec("nnunetv2") is None: + raise unittest.SkipTest("the trainer needs the nnunetv2 extra") + + +def flatten(transform) -> list[str]: + """Class names, descending into Compose and unwrapping RandomTransform.""" + if isinstance(transform, ComposeTransforms): + return [name for child in transform.transforms for name in flatten(child)] + wrapped = getattr(transform, "transform", None) + if type(transform).__name__ == "RandomTransform" and wrapped is not None: + return [type(wrapped).__name__] + return [type(transform).__name__] + + +def training_transforms(config_name: str, **overrides) -> list[str]: + from smauglab.trainers.nnUNetTrainerDAExt import nnUNetTrainerDAExtGPU + + os.environ["SMAUGLAB_PARAMS_JSON"] = str(CONFIGS / config_name) + kwargs = { + "patch_size": PATCH, + "rotation_for_DA": ROTATION, + "deep_supervision_scales": DS_SCALES, + "mirror_axes": (0, 1, 2), + "do_dummy_2d_data_aug": False, + "use_mask_for_norm": None, + "is_cascaded": False, + "foreground_labels": None, + "regions": None, + "ignore_label": None, + } + kwargs.update(overrides) + return flatten(nnUNetTrainerDAExtGPU.get_training_transforms(**kwargs)) + + +def cpu_block(config_name: str) -> list[str]: + """The CPU pipeline a config asks for, independent of the trainer.""" + config = load_config(str(CONFIGS / config_name)) + built = build_cpu_pipeline(config.section(Backend.CPU), do_dummy_2d_data_aug=False, patch_size=PATCH, rotation=ROTATION) + return [name for transform in built for name in flatten(transform)] + + +class TestOnlyOneTrainerRemains(unittest.TestCase): + def test_the_load_bearing_name_survives(self): + """nnU-Net writes the class name into every checkpoint and resolves the class + from it at inference, and several hundred trained runs record this one.""" + from smauglab.trainers import nnUNetTrainerDAExt + + self.assertTrue(hasattr(nnUNetTrainerDAExt, "nnUNetTrainerDAExtGPU")) + + def test_the_backend_specific_trainers_are_gone(self): + """They differed only in which config they defaulted to, which the config + itself now says. Neither had a single run on disk.""" + from smauglab.trainers import nnUNetTrainerDAExt + + for gone in ("nnUNetTrainerDAExtHybrid", "nnUNetTrainerDAExt"): + with self.subTest(trainer=gone): + self.assertFalse(hasattr(nnUNetTrainerDAExt, gone)) + + +class TestCompositionMatchesTheOldTrainers(unittest.TestCase): + """Each config must build what its dedicated trainer used to build.""" + + def test_gpu_config_matches_the_old_gpu_trainer(self): + self.assertEqual(training_transforms("transform_params_gpu.json"), ["SpatialTransform", "RemoveLabelTransform"]) + + def test_hybrid_config_matches_the_old_hybrid_trainer(self): + expected = [*cpu_block("transform_params_hybrid.json"), "RemoveLabelTransform"] + self.assertEqual(training_transforms("transform_params_hybrid.json"), expected) + + def test_cpu_config_builds_the_whole_cpu_pipeline(self): + got = training_transforms("transform_params.json") + self.assertEqual(got[: len(cpu_block("transform_params.json"))], cpu_block("transform_params.json")) + self.assertIn("SpatialTransform", got) + + def test_the_cpu_config_carries_the_spatial_transform_the_trainer_used_to_hardcode(self): + """The old CPU trainer appended a SpatialTransform with every probability at + 0 -- a no-op that only enforces the patch size. The merged trainer builds only + what the config names, so the config has to say it.""" + section = load_config(str(CONFIGS / "transform_params.json")).section(Backend.CPU) + self.assertIn("SpatialTransform", section) + spatial = section["SpatialTransform"] + self.assertEqual(spatial["p_rotation"], 0) + self.assertEqual(spatial["p_scaling"], 0) + self.assertEqual(spatial["p_elastic_deform"], 0) + self.assertEqual(spatial["mode_seg"], "nearest") + + +class TestDeepSupervisionPlacement(unittest.TestCase): + """Downsampling must follow whatever last deformed the mask. + + With GPU augmentations that is `train_step`, so it happens there; without them + nothing touches the mask after the dataloader and it belongs there. Getting this + backwards would train against targets that no longer match the image. + """ + + def test_a_gpu_config_leaves_downsampling_to_train_step(self): + for config in ("transform_params_gpu.json", "transform_params_hybrid.json"): + with self.subTest(config=config): + self.assertNotIn("DownsampleSegForDSTransform", training_transforms(config)) + + def test_a_cpu_only_config_downsamples_in_the_dataloader(self): + self.assertIn("DownsampleSegForDSTransform", training_transforms("transform_params.json")) + + def test_no_downsampling_when_no_scales_are_requested(self): + got = training_transforms("transform_params.json", deep_supervision_scales=None) + self.assertNotIn("DownsampleSegForDSTransform", got) + + +class TestDummy2D(unittest.TestCase): + def test_the_converters_bracket_the_spatial_transform(self): + got = training_transforms("transform_params_gpu.json", do_dummy_2d_data_aug=True) + spatial = got.index("SpatialTransform") + self.assertEqual(got[spatial - 1], "Convert3DTo2DTransform") + self.assertEqual(got[spatial + 1], "Convert2DTo3DTransform") + + +class TestConfigResolution(unittest.TestCase): + def setUp(self): + self._saved = {k: os.environ.get(k) for k in ("SMAUGLAB_PARAMS_JSON", "SMAUGLAB_PARAMS_GPU_JSON")} + for key in self._saved: + os.environ.pop(key, None) + + def tearDown(self): + for key, value in self._saved.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + def test_the_packaged_default_is_used_when_nothing_is_set(self): + from smauglab.trainers.nnUNetTrainerDAExt import resolve_config_path + + self.assertTrue(resolve_config_path().endswith("transform_params_gpu.json")) + + def test_the_new_variable_wins(self): + from smauglab.trainers.nnUNetTrainerDAExt import resolve_config_path + + os.environ["SMAUGLAB_PARAMS_JSON"] = "/new.json" + os.environ["SMAUGLAB_PARAMS_GPU_JSON"] = "/old.json" + self.assertEqual(resolve_config_path(), "/new.json") + + def test_the_old_variable_still_works_but_warns(self): + """segtransferaug/run_trainings.py sets it, and it drives every historical run.""" + from smauglab.trainers.nnUNetTrainerDAExt import resolve_config_path + + os.environ["SMAUGLAB_PARAMS_GPU_JSON"] = "/old.json" + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + self.assertEqual(resolve_config_path(), "/old.json") + self.assertTrue(any(issubclass(w.category, DeprecationWarning) for w in caught)) + + +class TestPipelineMode(unittest.TestCase): + def test_configs_default_to_the_sequential_pipeline(self): + self.assertIs(load_config(str(CONFIGS / "transform_params_gpu.json")).pipeline_mode(), PipelineMode.SEQUENTIAL) + + def test_the_list_configs_ask_for_the_random_order_pipeline(self): + """They carry a random_choose block, so they were written for the ChooseX + trainers; the config says that now instead of the class name.""" + for config in sorted(CONFIGS.glob("*-List*.json")): + with self.subTest(config=config.name): + self.assertIs(load_config(str(config)).pipeline_mode(), PipelineMode.RANDOM_ORDER) + + +if __name__ == "__main__": + unittest.main() From 216be9d7031cd280370f8bb449db5dba3b18974e Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 27 Aug 2026 12:56:44 +0000 Subject: [PATCH 2/2] fix: track the generated config, and point the CLI tests at a tracked one Both failures were the same thing: .gitignore has a repo-wide *.json, and these files were on the wrong side of it. smauglab/configs/all_augmentations.json is generated by `smauglab template --write` and checked by `smauglab template --check`, so it is source and has to be in the repo. Un-ignored and committed. Without it the staleness test could not read the file it was checking, and three CLI tests failed on a clean checkout while passing here. unit_tests/test_cli.py named transform_params_gpu_default01-23.json, which is one of the ~24 per-experiment configs the ignore rule deliberately keeps out. Pointed at transform_params_gpu.json, which is tracked. Naming an ignored file is the failure mode this whole class of bug keeps taking, so the constant says so. Co-Authored-By: Claude Opus 5 --- .gitignore | 3 + smauglab/configs/all_augmentations.json | 667 ++++++++++++++++++++++++ unit_tests/test_cli.py | 6 +- 3 files changed, 674 insertions(+), 2 deletions(-) create mode 100644 smauglab/configs/all_augmentations.json diff --git a/.gitignore b/.gitignore index 83fbfeb..3fdaf94 100644 --- a/.gitignore +++ b/.gitignore @@ -46,6 +46,9 @@ MANIFEST # migration table, the test fixtures and the migrator itself. !unit_tests/fixtures/**/*.json !migration/*.json +# Generated by `smauglab template --write` and checked by CI; it is source, +# not per-experiment output. +!smauglab/configs/all_augmentations.json *.yaml *.nii *.nii.gz diff --git a/smauglab/configs/all_augmentations.json b/smauglab/configs/all_augmentations.json new file mode 100644 index 0000000..b25b005 --- /dev/null +++ b/smauglab/configs/all_augmentations.json @@ -0,0 +1,667 @@ +{ + "GPU": { + "RandomFlipTransformGPU": { + "flip_axis": [ + 0 + ], + "keepdim": true, + "p": 1.0, + "p_batch": 1.0, + "same_on_batch": false + }, + "RandomAffineGPU": { + "align_corners": true, + "degrees": 10, + "keepdim": true, + "p": 0.5, + "p_batch": 1.0, + "resample": "BILINEAR", + "same_on_batch": false, + "scale": [ + 0.9, + 1.1 + ], + "shears": [ + -10, + 10, + -10, + 10, + -10, + 10 + ], + "translate": [ + 0.1, + 0.1, + 0.1 + ] + }, + "RandomSynthSegGPU": { + "apply_bias_field": true, + "apply_intensity_augmentation": true, + "apply_resolution": true, + "apply_to_channel": null, + "atlas_res": 1.0, + "bias_field_std": 0.7, + "bias_scale": 0.025, + "blur_range": 1.03, + "clip": 300.0, + "data_res": null, + "em_background_clusters_range": [ + 3, + 10 + ], + "em_background_label": 0, + "em_label_completion": false, + "em_max_fit_voxels": 100000, + "em_n_foreground_clusters": 2, + "em_n_iters": 20, + "em_same_on_batch": false, + "flip_axis": 2, + "gamma_std": 0.5, + "generation_classes": null, + "generation_labels": null, + "keepdim": true, + "max_res_aniso": 8.0, + "max_res_iso": 4.0, + "n_channels": 1, + "n_neutral_labels": null, + "nonlin_scale": 0.04, + "nonlin_std": 4.0, + "normalise": true, + "output_labels": null, + "p": 0.5, + "p_batch": 1.0, + "prior_distributions": "uniform", + "prior_means": null, + "prior_stds": null, + "randomise_res": true, + "rotation_bounds": 15.0, + "same_on_batch": false, + "scaling_bounds": 0.2, + "shearing_bounds": 0.012, + "svf_integration_steps": 7, + "thickness": null, + "translation_bounds": false + }, + "RandomPaletteGPU": { + "alpha_magnitude_range": [ + 0.5, + 2.0 + ], + "blur_sigmas": [ + 0.0, + 0.0, + 0.0, + 0.3, + 0.5, + 0.8 + ], + "c_choices": [ + 2, + 3, + 4, + 5, + 6 + ], + "dark_threshold": 0.01, + "keepdim": false, + "label_classes": null, + "label_remap_prob": 0.5, + "min_label_voxels": 4, + "n_kmeans_subsample": 10000, + "p": 1.0, + "p_batch": 1.0, + "s_choices": [ + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10 + ], + "same_on_batch": false, + "skip_parcellation_prob": 0.1, + "skip_sub_parc_prob": 0.4 + }, + "RandomDomainTransferGPU": { + "any_source": false, + "apply_to_channel": null, + "bank_path": null, + "bias_field_std": 0.0, + "bias_scale": 0.03, + "blend_concentration": 1.0, + "blend_targets": 1, + "include_self": true, + "keepdim": true, + "p": 0.2, + "p_batch": 1.0, + "p_class_mix": 0.0, + "p_spatial_mix": 0.0, + "pct": 1.0, + "same_on_batch": false, + "sigma": 2.0, + "source_label": null, + "spatial_mix_gain": 3.0, + "spatial_mix_scale": 0.03, + "targets": null, + "zscore_io": "auto" + }, + "RandomInverseGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomHistogramEqualizationGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomRedistributeSegGPU": { + "apply_to_channel": [ + 0 + ], + "dilation_iterations_range": [ + 1, + 3 + ], + "in_seg": 0.2, + "keepdim": true, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false, + "std_noise_range": [ + 0.1, + 0.3 + ] + }, + "RandomScharrGPU": { + "absolute": true, + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": true, + "same_on_batch": false + }, + "RandomLaplaceGPU": { + "absolute": false, + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomUnsharpMaskGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false, + "sigma": 1.0, + "unsharp_amount": 1.5 + }, + "RandomRandConvGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "kernel_sizes": [ + 1, + 3, + 5, + 7 + ], + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomClampGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "max_clamp_amount": 0.0, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomGaussianNoiseGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mean": 0.0, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "same_on_batch": false, + "std": 1.0 + }, + "RandomGaussianBlurGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "mix_prob": 0.0, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false, + "sigma": 1.0 + }, + "RandomBrightnessGPU": { + "apply_to_channel": [ + 0 + ], + "brightness_range": [ + 0.5, + 1.5 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "same_on_batch": false + }, + "RandomGammaGPU": { + "apply_to_channel": [ + 0 + ], + "gamma_range": [ + 0.7, + 1.5 + ], + "in_seg": 0.0, + "keepdim": false, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomInvGammaGPU": { + "apply_to_channel": [ + 0 + ], + "gamma_range": [ + 0.7, + 1.5 + ], + "in_seg": 0.0, + "keepdim": false, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomContrastGPU": { + "apply_to_channel": [ + 0 + ], + "contrast_range": [ + 0.75, + 1.25 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomLog1pGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomSqrtGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomSinGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomExpGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomSigmoidGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "mix_in_out": false, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "RandomLowResTransformGPU": { + "keepdim": true, + "p": 1.0, + "p_batch": 1.0, + "same_on_batch": false, + "scale": [ + 0.3, + 1.0 + ] + }, + "RandomAcqTransformGPU": { + "apply_to_channel": [ + 0 + ], + "keepdim": true, + "p": 1.0, + "p_batch": 1.0, + "same_on_batch": false, + "scale": [ + 0.3, + 1.0 + ] + }, + "RandomCropTransformGPU": { + "crop": [ + 1.0, + 1.0 + ], + "keepdim": true, + "p": 1.0, + "p_batch": 1.0, + "pos": [ + 0.0, + 1.0 + ], + "same_on_batch": false + }, + "RandomBiasFieldGPU": { + "apply_to_channel": [ + 0 + ], + "coefficients": 0.5, + "in_seg": 0.0, + "invert": false, + "keepdim": true, + "mix_in_out": false, + "order": 3, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "retain_stats": false, + "same_on_batch": false + }, + "ZscoreNormalizationGPU": { + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "keepdim": true, + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0 + } + }, + "CPU": { + "LaplaceConvTransform": { + "absolute": false, + "p": 1.0, + "retain_stats": false + }, + "ScharrConvTransform": { + "absolute": true, + "p": 1.0, + "retain_stats": false + }, + "Log1pTransform": { + "p": 1.0, + "retain_stats": false + }, + "SqrtTransform": { + "p": 1.0, + "retain_stats": false + }, + "SinTransform": { + "p": 1.0, + "retain_stats": false + }, + "ExpTransform": { + "p": 1.0, + "retain_stats": false + }, + "SigmoidTransform": { + "p": 1.0, + "retain_stats": false + }, + "HistogramEqualTransform": { + "p": 1.0, + "retain_stats": false + }, + "RedistributeTransform": { + "classes": null, + "in_seg": 0.2, + "p": 1.0, + "retain_stats": false + }, + "ShapeTransform": { + "ignore_axes": [], + "p": 1.0, + "shape_min": 1 + }, + "ArtifactTransform": { + "bias_field": false, + "blur": false, + "ghosting": false, + "motion": false, + "noise": false, + "p": 1.0, + "random_pick": false, + "spike": false, + "swap": false + }, + "SpatialCustomTransform": { + "affine": false, + "anisotropy": false, + "elastic": false, + "flip": false, + "p": 1.0, + "random_pick": false + }, + "SpatialTransform": { + "align_corners": false, + "bg_style_seg_sampling": true, + "border_mode_seg": "zeros", + "center_deformation": true, + "elastic_deform_magnitude": [ + 0, + 0.2 + ], + "elastic_deform_scale": [ + 0, + 0.2 + ], + "mode_image": "bilinear", + "mode_seg": "bilinear", + "p_elastic_deform": 0, + "p_rot_per_axis": 1, + "p_rotation": 0, + "p_scaling": 0, + "p_synchronize_def_scale_across_axes": 0, + "p_synchronize_scaling_across_axes": 0, + "padding_mode_image": "zeros", + "padding_value_image": 0, + "padding_value_seg": 0, + "patch_center_dist_from_border": null, + "random_crop": null, + "scaling": [ + 0.7, + 1.3 + ] + }, + "GaussianNoiseTransform": { + "noise_variance": [ + 0, + 0.1 + ], + "p": 1.0, + "p_per_channel": 1.0, + "synchronize_channels": false + }, + "GaussianBlurTransform": { + "benchmark": false, + "blur_sigma": [ + 1, + 5 + ], + "p": 1.0, + "p_per_channel": 1, + "synchronize_axes": false, + "synchronize_channels": false + }, + "MultiplicativeBrightnessTransform": { + "multiplier_range": null, + "p": 1.0, + "p_per_channel": 1, + "synchronize_channels": null + }, + "ContrastTransform": { + "contrast_range": null, + "p": 1.0, + "p_per_channel": 1.0, + "preserve_range": null, + "synchronize_channels": null + }, + "SimulateLowResolutionTransform": { + "allowed_channels": null, + "ignore_axes": null, + "p": 1.0, + "p_per_channel": 1, + "scale": null, + "synchronize_axes": null, + "synchronize_channels": null + }, + "InvertedGammaTransform": { + "gamma": [ + 0.7, + 1.5 + ], + "p": 1.0, + "p_per_channel": 1, + "p_retain_stats": 1, + "synchronize_channels": false + }, + "GammaTransform": { + "gamma": null, + "p": 1.0, + "p_invert_image": null, + "p_per_channel": null, + "p_retain_stats": null, + "synchronize_channels": null + }, + "MirrorTransform": { + "allowed_axes": null + }, + "ZscoreNormalization": { + "p": 1.0 + } + } +} diff --git a/unit_tests/test_cli.py b/unit_tests/test_cli.py index a3f134c..320977b 100644 --- a/unit_tests/test_cli.py +++ b/unit_tests/test_cli.py @@ -17,8 +17,10 @@ from smauglab import cli REPO = Path(__file__).resolve().parent.parent -DEFAULT_GPU = REPO / "smauglab" / "configs" / "transform_params_gpu_default01-23.json" -LEGACY = REPO / "unit_tests" / "fixtures" / "legacy_configs" / "configs" / "transform_params_gpu_default01-23.json" +# A *tracked* config. Most of smauglab/configs is gitignored (per-experiment +# sweeps), so naming one of those makes the test pass locally and fail in CI. +DEFAULT_GPU = REPO / "smauglab" / "configs" / "transform_params_gpu.json" +LEGACY = REPO / "unit_tests" / "fixtures" / "legacy_configs" / "configs" / "transform_params_gpu.json" def run(*argv: str) -> tuple[int, str]: