[Hendrik-code]

What

A new pipeline key, split, that makes groups of GPU transforms mutually exclusive:

"pipeline": {"split": {"paths": [["RandomSynthSegGPU"],
                                 ["RandomInverseGPU", "RandomHistogramEqualizationGPU",
                                  "RandomRedistributeSegGPU", "RandomScharrGPU",
                                  "RandomUnsharpMaskGPU", "RandomRandConvGPU"]],
                       "p": [0.5, 0.5]}}

Each batch takes exactly one path, path i with probability p[i]:

• the transforms of the paths not taken are skipped for that batch (no parameter draw, image and mask untouched);
• the transforms of the path taken run as usual, each still subject to its own p;
• every transform the split does not name runs on every batch.

Pipeline order is unchanged: a split decides which transforms run, never when, so a path does not have to be contiguous. An empty path is allowed and means "none of the others".

Why

The question was "SynthSeg on half the batches, our transfer augmentations on the other half, the general enhancements on all of them". The config surface could not say "not both": adding RandomSynthSegGPU at p: 0.5 to the paper config stacks it, because SynthSeg sits ahead of every TA in PIPELINE_ORDER, so on the batches it fires the TAs then run on the synthetic image.

With the split above and SynthSeg at p: 1.0, each path reproduces one of the two shipped paper configs exactly: transform_params_one-sequence-to-segment-them-all.json and transform_params_paper-synthseg.json differ only in the SynthSeg block and in the p of those six TAs.

How

• SmaugConfig.pipeline_split() returns a SplitSpec(paths, p). Validation (all problems reported together): at least two paths, one p per path, each in [0, 1], summing to 1; every name in the GPU section (with a did-you-mean) and in one path at most; sequential mode only, since the bucketing modes hide the members inside a RandomChooseXTransformsGPU; and not with same_on_batch: false, because the path is drawn once per batch.
• AugmentationSequentialCustom(split=...) resolves each path to its child stages by class name (a config key is the class name) and refuses a name that is not exactly one stage. It overrides get_forward_sequence: on the parameter-drawing call it draws the path and filters out the other paths' stages. forward and the mask pass iterate the parameters that call produced, so a skipped stage is skipped for image and mask alike, and replays with explicit params are untouched. The same hook feat: shuffle a group's order in place with pipeline.shuffle #143 used.
• The draw uses rng.shared_rand, so it follows torch.manual_seed and DDP ranks agree on the path.
• The nnU-Net trainer and AugTransformsGPU both pass it through. The trainer prints the split under its config line.

Without split nothing changes: the override returns kornia's sequence as is.

Tests

unit_tests/test_pipeline_split.py (21 tests): parsing and every rejection; no split keeps every stage; over 4000 seeded draws the two paths never co-occur, the shared stages always run and the path frequency is within 5 sd of p; an empty path; the draw follows the torch seed; through a real forward with record_applications, exactly one path runs per batch and a skipped stage leaves the mask alone; a name that is not a stage raises; AugTransformsGPU and the trainer (constructed with nnU-Net's base __init__ patched out) carry the split. Removing the trainer's split= argument fails the trainer test.

Full CPU suite: 494 passed, 13 skipped. Ruff clean (local 0.17.1; CI pins 0.16.1). mypy 2.1.0 clean on smauglab/ (CI pins 2.4.0).

An nnU-Net training with the config above (nnUNetTrainerDAExtGPU, fold 0, 1000 epochs) started on Dataset803 on 2026-11-13, with Dataset804 queued behind it. The trainer logs the split at startup and runs at ~63 s/epoch on an A40, against ~64 s for the stacked SynthSeg p=0.5 arm. A pre-flight on 4000 parameter draws of that config took the paths 1973/2027 times, never both on one draw, with each TA at its paper p on its own path.

🤖 Generated with Claude Code