[Hendrik-code]

What does this change?

Adds an optional second loss term to nnUNetTrainerDAExtGPU: the augmented patch and the unaugmented one both go through the network, and their predictions are compared with a soft Dice. Off unless a config asks for it.

{
  "loss": { "consistency": { "enabled": true, "weight": 0.5, "detach_reference": true } },
  "GPU": { "...": {} }
}

The reference is taken after the geometric transforms and before the intensity ones. The geometric group is the only one that moves voxels, so running it once and sharing the output is what makes the two images voxel-aligned and lets a single segmentation be correct for both branches.

Mechanics:

• build_gpu_stages(..., split_shared=True) returns (shared, augmented). build_gpu_pipeline is now a thin wrapper over it (shared + augmented), so the unsplit path cannot drift. _split_shared raises rather than hoisting anything that is not already a prefix of the pipeline.
• New AugType.OCC, holding RandomCropTransformGPU. build.SEQUENTIAL_GROUPS lists it beside GEO so no built pipeline changes; build.SHARED_GROUPS leaves it out of the shared stage.
• New top-level loss config section, validated like pipeline with difflib suggestions. It is inside config_hash, so enabling the term gives a new ...-c-<hash> run folder, and _record_config already copies it into the run as provenance.
• Reuses nnU-Net's MemoryEfficientSoftDiceLoss with the trainer's own batch_dice / smooth / region settings, so the two Dice terms in the sum are computed identically.

Why?

Nothing in the current objective expresses the invariance the augmentations are there to produce. This states it directly as a loss.

Why the crop is not GEO

RandomCropTransformGPU occludes the field of view rather than translating contents, and PIPELINE_ORDER puts it near the end for a reason: run in front of the intensity stage, the flat fill it leaves behind gets texture painted into it — the failure its own pad_value docstring describes for PALETTE and the edge filters. A new AugType rather than an exclusion list because the group is what _buckets already reads, and because GEO now means exactly "moves voxels", which is the property the split depends on.

Why not SoftDiceLoss

On nnunetv2 2.6, get_tp_fp_fn_tn builds y_onehot as a bool tensor and computes fp = net_output * (~y_onehot). A soft target of matching shape takes the y_onehot = gt branch and ~gt raises TypeError: ~ (operator.invert) is only implemented on integer and Boolean-type tensors. MemoryEfficientSoftDiceLoss, from the same module and the one _build_loss itself uses, only ever does y_onehot.sum(axes) and x * y_onehot, so it takes the soft target directly. (2.4.2 used 1 - y_onehot and would have worked — the kind of upstream drift pyproject.toml already documents.)

How was it tested?

• pytest passes locally
• pre-commit run --all-files passes
• Added tests
• Ran a training job end to end

Unit tests — 50 new tests, 190 new subtests.

where	what
test_consistency_loss.py (new)	The term's arithmetic against a longhand soft Dice; that only the augmented branch gets gradients when detached; that the deep-supervision heads never reach it; the epoch summary; the verdict=consistency NaN report; that chaining the two stages reproduces the single container exactly on the shipped config; the guardrails
test_builder.py	That the two stages concatenate to the unsplit pipeline in every mode, that the crop never enters the shared stage, that the prefix rule is enforced, and a sha256 over all 84 (config, mode) pipeline structures as the regression guard
test_config_schema.py	The loss section: defaults, copies, difflib suggestions, type and range checks, and that a JSON null is rejected rather than read as "use the default"
test_registry.py, test_cli.py	The new group's membership and that smauglab list --group occ works

The pipeline-structure hash was captured on main before any edit and is unchanged after, so nothing about the existing pipelines moved.

Two environments, because this machine's default interpreter has nnunetv2 2.4.2:

env	nnunetv2	before	after
seg11	2.4.2	9 failed, 456 passed	9 failed, 506 passed
nnunet11	2.6.2	2 failed (targeted)	2 failed (targeted)

The 9 are the documented >=2.5 signature mismatch; the 2 are a batchgeneratorsv2 RemoveLabelTansform → RemoveLabelTransform rename. Both sets are identical before and after this change. mypy smauglab/ clean (locally 2.1.0; the pin is 2.4.0, so CI is the real check).

End to end, real trainer and dataloader, Dataset400_nose_nosemain 3d_fullres fold 0, one GPU:

[term off ] losses=[2.38409, 2.33402, 2.25938]  supervised=[2.38409, 2.33402, 2.25938]
[term on  ] losses=[2.31166, 2.27615, 2.19719]  supervised=[2.38999, 2.33600, 2.25735]
[term on  ] log: consistency: soft dice 0.1322 over 3 batches, weight 0.5, contributing -0.0661 to train_loss
[weight 0 ] losses=[2.38443, 2.32928, 2.26386]  supervised=[2.38443, 2.32928, 2.26386]
[weight 0 ] log: consistency: soft dice 0.1317 over 3 batches, weight 0.0, contributing -0.0000 to train_loss
shared geometric stage: 2 transforms, intensity stage: 24 transforms

Total equals supervised plus weight times the term on every step, and weight: 0 measures and logs the term while leaving the loss bit-identical to the supervised value. A same-config control established that this setup is not reproducible run to run (the augmented batch differs by up to 1e1 between two identical runs), so an on-vs-off batch comparison would have said nothing; bit-identity of the split is pinned deterministically on CPU in the unit tests instead.

Anything reviewers should look at closely?

1. The crop interaction, which is not masked. The crop is not applied to the reference, so in the ~25% of batches where it fires the term asks the network to agree with the reference inside the occluded region, where the supervised target says background. The two terms pull against each other there. The trainer warns at startup naming the actual p, and the README says to set p: 0. Masking the term to the kept box is possible — a ones-channel through the crop's mask path gives the exact indicator — but it was deliberately left out of this PR.
2. Only the full-resolution head takes part. The lower heads carry nnU-Net's own _build_loss weights, including weights[-1] = 0 and its DDP-dependent 1e-6 variant; giving the consistency term a second weighting scheme would duplicate a decision that belongs upstream.
3. No .clone() of the reference. It is an alias of the geometric stage's output, justified by test_no_inplace_mutation.py plus a new test that runs the shipped configuration rather than default-constructed transforms.
4. train_loss is no longer comparable across the switch, since it now carries both terms. Hence the separate epoch line. nnUNetLogger.log asserts the key already exists in my_fantastic_logging, so a proper curve would need an upstream patch.
5. With a soft target the term does not bottom out at -1 when the branches agree; its minimiser is the reference's hard argmax, so it sharpens as well as aligns. Pinned by a test so it is not later "fixed".
6. One semantic change beyond the loss: registry.names(Backend.GPU, group=AugType.GEO), filter_by_group(..., AugType.GEO) and smauglab list --group geo no longer return RandomCropTransformGPU; it is under occ. In random_order mode a crop-enabled config would also now follow the shuffled buckets instead of leading — no shipped config combines the two (the three *-List*.json use random_order and have no crop; all four crop configs are sequential), so no published run is affected.

Checklist

• Augmentation behaviour is unchanged (proved by the 84-pipeline hash) — the only behaviour change is opt-in via loss.consistency
• README augmentation matrix regenerated (smauglab matrix --write; one cell, crop | GEO → crop | OCC). all_augmentations.json is unchanged, as render_template does not read the group
• New config section documented in the README with its keys and caveats
• CONTRIBUTING.md test table gained a row

🤖 Generated with Claude Code