Skip to content

refactor: extract the nnU-Net transform tail shared by five trainers - #64

Open
Hendrik-code wants to merge 1 commit into
hm/dedup-tailsfrom
hm/dedup-trainer-tail
Open

refactor: extract the nnU-Net transform tail shared by five trainers#64
Hendrik-code wants to merge 1 commit into
hm/dedup-tailsfrom
hm/dedup-trainer-tail

Conversation

@Hendrik-code

Copy link
Copy Markdown
Collaborator

Stacked on #63.

Five get_training_transforms methods across nnUNetTrainerDAExt.py and nnUNetTrainerTest.py ended with a character-identical copy of the same nnU-Net sequence — intensity masking, the -1 label removal, the two cascade transforms, region conversion, deep-supervision downsampling. That's trainers/utils.py::nnunet_tail_transforms now. −255 / +134.

The one thing the copies differed in

Whether the DownsampleSegForDSTransform block was live or commented out:

# NOTE: DownsampleSegForDSTransform is now handled in train_step for GPU augmentations
# if deep_supervision_scales is not None:
#     transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales))

The GPU trainers carry it commented because 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 in train_step. That's now expressed by passing deep_supervision_scales=None rather than by a comment.

What is deliberately left alone

get_validation_transforms is not folded in. Its cascade branch adds only MoveSegAsOneHotToDataTransform, without the two RandomTransform wrappers — so it's a different sequence, not a sixth copy. Forcing it through the helper would need a flag that changes that branch, costing more than it saves.

How I know it changes nothing

I enumerated the transform list every trainer builds across 160 argument combinations — dummy-2D on/off × use_mask_for_norm × cascade × regions × deep supervision, for all five trainers — and compared the type and repr of every entry before and after:

IDENTICAL: all 160 trainer transform lists unchanged (addresses normalised)

(The only raw differences were <function ... at 0x...> addresses.)

That check is local-only and cannot run in CI. nnunetv2 is an optional extra and pip install -e ".[dev]" does not pull it, so nothing under trainers/ is imported by the test suite at all — CI will go green here without exercising a single line of this diff. That is exactly why this is a separate PR from #63, which CI does cover.

Full suite: 105 passed, 200 subtests (unchanged — none of them touch the trainers). ruff and mypy clean.

Five get_training_transforms methods across nnUNetTrainerDAExt.py and
nnUNetTrainerTest.py ended with a character-identical copy of the same nnU-Net
sequence: intensity masking, the -1 label removal, the two cascade transforms,
region conversion and deep-supervision downsampling. That is
trainers/utils.py::nnunet_tail_transforms now. -255 lines, +134.

The four full copies differed in exactly one thing: whether the
DownsampleSegForDSTransform block was live or commented out. The GPU trainers
carry it commented, 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. That is expressed by passing
deep_supervision_scales=None rather than by a comment.

get_validation_transforms is deliberately left alone. Its cascade branch adds
only MoveSegAsOneHotToDataTransform, without the two RandomTransform wrappers, so
it is a different sequence rather than a sixth copy of this one. Forcing it
through the helper would need a flag that changes that branch, which costs more
than it saves.

Verified by enumerating the transform list every trainer builds across 160
argument combinations -- dummy 2D on/off, mask, cascade, regions, deep
supervision -- and comparing the type and repr of each entry before and after.
All 160 are identical once function memory addresses are normalised out.

That check is local-only and not repeatable in CI: nnunetv2 is an optional extra
and `pip install -e ".[dev]"` does not pull it, so nothing under trainers/ is
imported by the test suite at all. It is the reason this is a separate commit
from the gpu/contrast.py extraction, which CI does cover.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Copilot AI lite review requested due to automatic review settings August 20, 2026 12:42

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants