diff --git a/smauglab/config.py b/smauglab/config.py index 93bd7ca..b12fe43 100644 --- a/smauglab/config.py +++ b/smauglab/config.py @@ -30,7 +30,7 @@ from typing import Any from smauglab import registry -from smauglab.registry import Backend, InvalidConfigError +from smauglab.registry import AugType, Backend, InvalidConfigError class PipelineMode(str, Enum): @@ -96,7 +96,11 @@ def validate_section(section: dict, backend: Backend, *, source: str = " RESERVED_SECTIONS = ("pipeline",) #: Keys the `pipeline` section may hold. -PIPELINE_KEYS = ("mode", "order", "random_choose", "same_on_batch") +PIPELINE_KEYS = ("mode", "order", "random_choose", "same_on_batch", "shuffle") + +#: Groups `pipeline.shuffle` may name. GEO is absent: geometry moves the mask, and the +#: transforms after it read the segmentation it leaves behind. +SHUFFLE_GROUPS = (AugType.TA, AugType.GE) #: What to do about a config that has no backend section. Quoted by both errors that #: can report it, so the two cannot drift apart. @@ -176,6 +180,25 @@ def same_on_batch(self) -> bool: raw = self.payload.get("pipeline", {}).get("same_on_batch") return True if raw is None else bool(raw) + def shuffle_groups(self) -> tuple[AugType, ...]: + """The augmentation groups whose order is redrawn on every batch. + + `"pipeline": {"shuffle": ["TA"]}` permutes the transfer augmentations among the + positions they already hold in the sequential pipeline, and leaves everything + else where it is: the general enhancements between them, geometry in front, and + each transform's own `p`. `["GE"]` does the same for the general enhancements, + and `["TA", "GE"]` shuffles both, each among its own positions. Transforms + marked `force_sequential` keep their place. + + The `random_order` modes are not this. They also reorder, but they pull each + group into a `RandomChooseXTransformsGPU` bucket, which moves every TA ahead of + every GE and calls `apply_transform` directly instead of going through kornia. + A shuffled sequential pipeline differs from the unshuffled one in the order + alone. Empty by default, which is the pipeline every published run trained on. + """ + raw = self.payload.get("pipeline", {}).get("shuffle") or [] + return tuple(AugType(g) for g in raw) + def names(self, backend: Backend) -> list[str]: return [k for k in self.section(backend) if not k.startswith("_")] @@ -246,6 +269,17 @@ def _pipeline_problems(self) -> list[str]: close = difflib.get_close_matches(str(mode), valid, n=2, cutoff=0.5) hint = f" Did you mean: {', '.join(close)}?" if close else "" problems.append(f"pipeline.mode: unknown mode {mode!r}. Accepted: {', '.join(valid)}.{hint}") + + shuffle = pipeline.get("shuffle") + if shuffle is not None: + valid_groups = [g.value for g in SHUFFLE_GROUPS] + if not isinstance(shuffle, list) or any(g not in valid_groups for g in shuffle) or len(set(shuffle)) != len(shuffle): + problems.append(f"pipeline.shuffle: expected a list of distinct groups from {', '.join(valid_groups)}, got {shuffle!r}") + elif shuffle and mode not in (None, PipelineMode.SEQUENTIAL.value): + problems.append( + f"pipeline.shuffle: only applies to the sequential mode; {mode!r} already orders its groups through " + "random_choose (ta_random_order / ge_random_order)" + ) return problems diff --git a/smauglab/trainers/nnUNetTrainerDAExt.py b/smauglab/trainers/nnUNetTrainerDAExt.py index a9dac85..bf91b31 100644 --- a/smauglab/trainers/nnUNetTrainerDAExt.py +++ b/smauglab/trainers/nnUNetTrainerDAExt.py @@ -40,7 +40,7 @@ from smauglab.config import load_config from smauglab.registry import Backend from smauglab.trainers.utils import DownsampleSegForDSTransformCustom, nnunet_tail_transforms -from smauglab.transforms.build import build_cpu_pipeline, build_gpu_pipeline +from smauglab.transforms.build import build_cpu_pipeline, build_gpu_pipeline, shuffle_slots from smauglab.transforms.gpu.base import AugmentationSequentialCustom #: Env var naming the config to train with. @@ -187,23 +187,28 @@ def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dic # CPU-only config costs nothing per step and train_step stays a no-op. self.transforms: AugmentationSequentialCustom | None = None if _has_gpu_augmentations(config): + gpu_pipeline = build_gpu_pipeline( + config.section(Backend.GPU), + mode=config.pipeline_mode(), + options=config.pipeline_options("random_choose"), + source=config.source, + order_source=config.order_source(), + ) self.transforms = AugmentationSequentialCustom( - *build_gpu_pipeline( - config.section(Backend.GPU), - mode=config.pipeline_mode(), - options=config.pipeline_options("random_choose"), - source=config.source, - order_source=config.order_source(), - ), + *gpu_pipeline, data_keys=["input", "mask"], # See AugTransformsGPU: True forces the flag onto every child and # spends each transform's `p` once for the whole batch. same_on_batch=True if config.same_on_batch() else None, + shuffle_slots=shuffle_slots(gpu_pipeline, config.shuffle_groups()), ).to(self.device) print(f"Using SmaugLab transforms from: {json_path}") + groups = config.shuffle_groups() + shuffle = f", shuffle: {'+'.join(g.value for g in groups)}" if groups else "" print( - f" CPU: {len(config.names(Backend.CPU))} augmentations, GPU: {len(config.names(Backend.GPU))}, mode: {config.pipeline_mode().value}" + f" CPU: {len(config.names(Backend.CPU))} augmentations, GPU: {len(config.names(Backend.GPU))}, " + f"mode: {config.pipeline_mode().value}{shuffle}" ) _record_config(json_path, os.path.join(self.output_folder, "transform_params_used_for_training.json")) diff --git a/smauglab/transforms/build.py b/smauglab/transforms/build.py index 63fe621..00bf558 100644 --- a/smauglab/transforms/build.py +++ b/smauglab/transforms/build.py @@ -14,6 +14,7 @@ from __future__ import annotations +from collections.abc import Sequence from typing import Any from smauglab import registry @@ -142,6 +143,21 @@ def build_gpu_pipeline( return out +def shuffle_slots(pipeline: Sequence[Any], groups: Sequence[AugType]) -> list[list[int]]: + """Per group, the positions in a sequential GPU `pipeline` that `pipeline.shuffle` permutes. + + Read back from the registry by class name, which the registry guarantees is the + config key. `force_sequential` transforms keep their place, as they do in the + `random_order` buckets. + """ + if not groups: + # Returned before the lookup: a bucket-mode pipeline holds RandomChooseX + # containers, which are not registered. + return [] + entries = [registry.get(type(t).__name__, Backend.GPU) for t in pipeline] + return [[i for i, e in enumerate(entries) if e.group is group and not e.force_sequential] for group in groups] + + def build_cpu_pipeline( section: dict, *, diff --git a/smauglab/transforms/gpu/base.py b/smauglab/transforms/gpu/base.py index 65490ad..0ea9e5d 100644 --- a/smauglab/transforms/gpu/base.py +++ b/smauglab/transforms/gpu/base.py @@ -2,7 +2,7 @@ import copy import functools import warnings -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from typing import Any, Protocol, Union import kornia.augmentation as K @@ -150,9 +150,13 @@ def __init__( random_apply_weights: list[float] | None = None, transformation_matrix_mode: str = "silent", extra_args: dict[DataKey, dict[str, Any]] | None = None, + shuffle_slots: Sequence[Sequence[int]] = (), ) -> None: self._transform_matrix: Tensor | None self._transform_matrices: list[Tensor | None] = [] + # Positions whose occupants are permuted among themselves on each forward; see + # `get_forward_sequence`. A group of one has nothing to swap with. + self._shuffle_slots = tuple(tuple(slots) for slots in shuffle_slots if len(slots) > 1) super().__init__( *args, @@ -198,6 +202,23 @@ def __init__( self._transform_matrix = None self.extra_args = extra_args or {DataKey.MASK: {"resample": Resample.NEAREST, "align_corners": None}} + def get_forward_sequence(self, params: list[ParamItem] | None = None) -> Iterator[tuple[str, Module]]: + """kornia's sequence, with each `shuffle_slots` group permuted among its own positions. + + Drawn once per call, which is once per batch, so every sample of a batch shares + one order even when `same_on_batch` is off. A replay (`params` given) follows the + recorded order, and the mask is carried through the same `params` as the image. + """ + sequence = super().get_forward_sequence(params) + if params is not None or not self._shuffle_slots: + return sequence + children = list(sequence) + order = list(range(len(children))) + for slots in self._shuffle_slots: + for slot, k in zip(slots, torch.randperm(len(slots)).tolist()): + order[slot] = slots[k] + return iter([children[i] for i in order]) + def forward_parameters(self, batch_shape: torch.Size) -> list[ParamItem]: """kornia's own loop, with each leaf's draw routed through `_leaf_parameters`. diff --git a/smauglab/transforms/gpu/transforms.py b/smauglab/transforms/gpu/transforms.py index c57ed3c..53f8db9 100644 --- a/smauglab/transforms/gpu/transforms.py +++ b/smauglab/transforms/gpu/transforms.py @@ -1,6 +1,6 @@ from smauglab.config import PipelineMode, load_config from smauglab.registry import Backend -from smauglab.transforms.build import build_gpu_pipeline +from smauglab.transforms.build import build_gpu_pipeline, shuffle_slots from smauglab.transforms.gpu.base import AugmentationSequentialCustom @@ -36,4 +36,5 @@ def __init__(self, json_path: str): *transforms, data_keys=["input", "mask"], same_on_batch=True if config.same_on_batch() else None, + shuffle_slots=shuffle_slots(transforms, config.shuffle_groups()), ) diff --git a/unit_tests/test_pipeline_shuffle.py b/unit_tests/test_pipeline_shuffle.py new file mode 100644 index 0000000..e93b79e --- /dev/null +++ b/unit_tests/test_pipeline_shuffle.py @@ -0,0 +1,146 @@ +"""`pipeline.shuffle`: the order of a group is redrawn per batch, and nothing else moves.""" + +from __future__ import annotations + +import json +import tempfile +from pathlib import Path + +import torch + +from smauglab.config import SmaugConfig, load_config +from smauglab.registry import AugType, Backend, InvalidConfigError +from smauglab.transforms.build import build_gpu_pipeline, shuffle_slots +from smauglab.transforms.gpu.base import AugmentationSequentialCustom, record_applications +from smauglab.transforms.gpu.transforms import AugTransformsGPU +from unit_tests.helpers import SmaugLabTestCase, configs_dir + +PAPER = "transform_params_one-sequence-to-segment-them-all.json" + +#: Three of each group at p=1, so every draw applies all of them and only the order varies. +SECTION = { + "RandomFlipTransformGPU": {"flip_axis": [0], "p": 1.0}, + "RandomInverseGPU": {"p": 1.0}, + "RandomGaussianNoiseGPU": {"p": 1.0}, + "RandomScharrGPU": {"p": 1.0}, + "RandomBrightnessGPU": {"p": 1.0}, + "RandomGammaGPU": {"p": 1.0}, + "RandomLowResTransformGPU": {"p": 1.0}, + "RandomBiasFieldGPU": {"p": 1.0}, +} + + +def names(pipeline) -> list[str]: + return [type(module).__name__ for _, module in pipeline.get_forward_sequence()] + + +class TestConfig(SmaugLabTestCase): + def test_the_default_shuffles_nothing(self): + self.assertEqual(SmaugConfig({"GPU": dict(SECTION)}).shuffle_groups(), ()) + + def test_groups_are_read_in_order(self): + config = SmaugConfig({"GPU": dict(SECTION), "pipeline": {"shuffle": ["GE", "TA"]}}) + self.assertEqual(config.shuffle_groups(), (AugType.GE, AugType.TA)) + + def test_bad_values_are_rejected(self): + for value in (["GEO"], ["TA", "TA"], "TA", ["ta"]): + with self.subTest(shuffle=value), self.assertRaises(InvalidConfigError): + SmaugConfig({"GPU": dict(SECTION), "pipeline": {"shuffle": value}}) + + def test_the_bucket_modes_reject_it(self): + with self.assertRaises(InvalidConfigError) as caught: + SmaugConfig({"GPU": dict(SECTION), "pipeline": {"mode": "random_order", "shuffle": ["TA"]}}) + self.assertIn("random_choose", str(caught.exception)) + + +class TestSlots(SmaugLabTestCase): + def test_paper_slots_are_the_group_positions(self): + config = load_config(str(configs_dir() / PAPER)) + pipeline = build_gpu_pipeline(config.section(Backend.GPU)) + ta, ge = shuffle_slots(pipeline, (AugType.TA, AugType.GE)) + at = {type(t).__name__: i for i, t in enumerate(pipeline)} + + self.assertEqual( + [type(pipeline[i]).__name__ for i in ta], + [ + "RandomInverseGPU", + "RandomHistogramEqualizationGPU", + "RandomRedistributeSegGPU", + "RandomScharrGPU", + "RandomUnsharpMaskGPU", + "RandomRandConvGPU", + "RandomLog1pGPU", + "RandomSqrtGPU", + "RandomSinGPU", + "RandomExpGPU", + "RandomSigmoidGPU", + "RandomBiasFieldGPU", + ], + ) + self.assertNotIn(at["RandomLowResTransformGPU"], ge, "force_sequential must keep its place") + self.assertNotIn(at["RandomFlipTransformGPU"], ta + ge) + self.assertIn(at["RandomGammaGPU"], ge) + + +class TestOrder(SmaugLabTestCase): + def _pipeline(self, groups): + built = build_gpu_pipeline(dict(SECTION)) + return AugmentationSequentialCustom( + *built, data_keys=["input", "mask"], same_on_batch=True, shuffle_slots=shuffle_slots(built, groups) + ) + + def test_without_groups_the_order_is_the_pipeline_order(self): + pipeline = self._pipeline(()) + fixed = [type(t).__name__ for t in pipeline.children()] + for _ in range(5): + self.assertEqual(names(pipeline), fixed) + + def test_only_the_named_group_moves_and_only_among_its_own_slots(self): + pipeline = self._pipeline((AugType.TA,)) + fixed = [type(t).__name__ for t in pipeline.children()] + ta = {"RandomInverseGPU", "RandomScharrGPU", "RandomBiasFieldGPU"} + slots = [i for i, n in enumerate(fixed) if n in ta] + + seen = set() + for _ in range(60): + order = names(pipeline) + self.assertEqual([n for i, n in enumerate(order) if i not in slots], [n for i, n in enumerate(fixed) if i not in slots]) + self.assertEqual({order[i] for i in slots}, ta) + seen.add(tuple(order[i] for i in slots)) + self.assertEqual(len(seen), 6, "every permutation of three should turn up in 60 draws") + + def test_both_groups_move_independently(self): + pipeline = self._pipeline((AugType.TA, AugType.GE)) + fixed = [type(t).__name__ for t in pipeline.children()] + ge = {"RandomGaussianNoiseGPU", "RandomBrightnessGPU", "RandomGammaGPU"} + self.assertEqual(names(pipeline).index("RandomLowResTransformGPU"), fixed.index("RandomLowResTransformGPU")) + seen = {tuple(n for n in names(pipeline) if n in ge) for _ in range(60)} + self.assertEqual(len(seen), 6) + + def test_the_applied_order_is_the_drawn_order_and_the_mask_is_untouched(self): + pipeline = self._pipeline((AugType.TA, AugType.GE)) + image, seg = self.tiny_volume(), self.tiny_seg() + orders = set() + for _ in range(10): + with record_applications(pipeline) as fired: + _, mask = pipeline(image.clone(), seg.clone()) + applied = [e["transform"] for e in fired if e["target"] == "image"] + self.assertEqual(applied, [p.name.rsplit("_", 1)[0] for p in pipeline._params]) + # The flip is the only transform that moves the mask, and this one is symmetric. + self.assertTrue(torch.equal(mask, seg)) + orders.add(tuple(applied)) + self.assertGreater(len(orders), 1) + + +class TestDefaultUnchanged(SmaugLabTestCase): + def test_an_empty_shuffle_is_the_unshuffled_pipeline(self): + payload = json.loads((configs_dir() / PAPER).read_text()) + image, seg = self.tiny_volume(), self.tiny_seg() + outputs = [] + with tempfile.TemporaryDirectory() as tmp: + for i, pipeline_section in enumerate(({}, {"shuffle": []})): + path = Path(tmp) / f"transform_params_shuffle_{i}.json" + path.write_text(json.dumps({**payload, "pipeline": pipeline_section})) + torch.manual_seed(0) + outputs.append(AugTransformsGPU(json_path=str(path))(image.clone(), seg.clone())[0]) + self.assertTrue(torch.equal(*outputs))