feat: shuffle a group's order in place with pipeline.shuffle

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 = "<config>
 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
 
 

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"))

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,
     *,

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`.
 

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()),
         )

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))