Fix determinism in augmentation #13
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Hi MIC-DKFZ team,
This PR fixes an issue that was preventing fully deterministic behavior in the augmentation pipeline, even when a seed was provided. I noticed this behavior when trying to train a nnunet model deterministically, I'll submit a PR for the nnunet side too, hoping that can help !
Problem
The main issue was that several transforms used PyTorch's random functions (
torch.rand,torch.normal). In a multi-process environment likeMultiThreadedAugmenter(like in nnunet), the numpy RNG is correctly seeded in each worker, but the PyTorch RNG is not. This led to unpredictable results from any transform usingtorchfor randomness.I also found a couple of related issues: the
benchmark=Trueparameter inGaussianBlurTransformis inherently non-deterministic, andSpatialTransformhad some unstable randomness from mixing torch and numpy operations.Solution
The fix was to go through the library and make sure all random operations rely on NumPy's random generator. This way, everything is controlled by the single RNG that's properly seeded in the data loader's workers.
Here are the transforms that were updated:
RandomTransformSpatialTransform(for elastic deform)MirrorTransformSimulateLowResolutionTransformGaussianBlurTransformGaussianNoiseTransformMultiplicativeBrightnessTransformContrastTransformGammaTransformRicianNoiseTransformInvertImageTransformRemoveRandomConnectedComponentFromOneHotEncodingTransformApplyRandomBinaryOperatorTransformHow It's Tested
To make sure these fixes work and to catch any future regressions, I've added a new testing script,
determinism_test_pipeline.py.The script tests every transform in the library for both 2D and 3D data. The key part of the test is how it checks for determinism: it runs each transform twice, but for the second run, it only re-seeds the NumPy RNG. This aims to mimic the multi-worker environment. The script demonstrates the non determinism on the original code.
With these changes, the whole library now passes this test, so I believe we can be confident that augmentation pipelines are fully reproducible. This change should allow for fully reproducible training pipelines, which is a big deal for research. The performance impact should be minimal, and might even be a little better since some inefficient operations and benchmarking overhead were removed.
Thanks for maintaining this great library. Hope this helps, and let me know what you think!
Post Scriptum
Original image :
Augmented image :
Augmented image :
