Batched transforms#282
Merged
Merged
Conversation
* Add batched transforms for MAE training - BatchedRandFlipd - BatchedRandSharpend - BatchedRandLocalPixelShufflingd - BatchedRandHistogramShiftd - BatchedRandZStackShiftd - BatchedRand3DElasticd * Revert existing files to original state Keep only new transform files and necessary imports * Fix import sorting in new transform files * Fix import formatting in __init__.py * Format code with ruff formatter * bugfix for device * ruff format fix
edyoshikun
reviewed
Sep 16, 2025
edyoshikun
reviewed
Sep 16, 2025
edyoshikun
reviewed
Sep 16, 2025
edyoshikun
reviewed
Sep 16, 2025
edyoshikun
reviewed
Sep 16, 2025
edyoshikun
reviewed
Sep 17, 2025
edyoshikun
left a comment
edyoshikun
left a comment
Member
There was a problem hiding this comment.
we dont have tests for some of the transforms (ie 3d elastic, histogram shift, local pixel shuffle, sharpen, and z_stack shift).
Isn't the z_stack shift just an affine?
edyoshikun
reviewed
Sep 17, 2025
edyoshikun
left a comment
edyoshikun
left a comment
Member
There was a problem hiding this comment.
I couldn't find a better way to optimize the funciotns and I really liked the tests.
Contributor
Author
These were contributions from #284.
Should be numerically equivalent to a translation. But if this is always aligned in index space then it avoids resampling, which would be faster than affine. |
edyoshikun
approved these changes
Sep 17, 2025
edyoshikun
added a commit
that referenced
this pull request
Sep 19, 2025
* add test for scale intensity range percentiles transform * randomize test data * add decollate * detect leading slashes * configurable label column * fix typo * configurable example shape * type hint * fix example array * format * auto-align validation fov names * Adding some batched transforms for MAE training (#284) * Add batched transforms for MAE training - BatchedRandFlipd - BatchedRandSharpend - BatchedRandLocalPixelShufflingd - BatchedRandHistogramShiftd - BatchedRandZStackShiftd - BatchedRand3DElasticd * Revert existing files to original state Keep only new transform files and necessary imports * Fix import sorting in new transform files * Fix import formatting in __init__.py * Format code with ruff formatter * bugfix for device * ruff format fix * wip: random crop * format * use gather operation for cropping * register batched random crop * use unfold and bench * add dict version for random spatial crop * lint * benchmark crop * rename flip file * use tensorstore to stack * rename in register * add notebook to showcase transforms * batched center crop * batched gaussian noise * test and add to notebook * fix docstring * test random flip * add map version of batched gaussian noise * fix noise scaling and shifting * per-sample random flip * fix gaussian blur * rename blur to smooth to match monai * update smoothing tests * sigma-based truncation * print timer results * random adjust contrast * wrap monai * random scale intensity * update tests * allow final crop override * add center crop to example * remove redundant call method * check gamma value * use batched center crop * fix import * wip: fix gathering * limit worker to 1 * wip: execute augmentations in data hook * use the full loop for profiling * wip: update concat wrapper * fix negative sampling * update tests * revert to use threads * cache tensorstore arrays * update default num_workers * skip transforms for example array * relax flaky randomness test * explicitly configure cache pool * use new iohub API * only check cache when opening * fix transform initialization * remove unecessary tipletdatamodule inits * explain num_workers in the docstring * updating paths for profiling * fixed batch randintensityscale in the notebook --------- Co-authored-by: Ritvik <40371793+ritvikvasan@users.noreply.github.com> Co-authored-by: Eduardo Hirata-Miyasaki <edhiratam@gmail.com>
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
3 participants
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.
Add more batched transforms.
Tested (now includes all DynaCLR transforms, along with #276):
See an illustration of these in the new example notebook.
Not tested (#284):
Also updated the triplet dataset and datamodule for end-to-end batching. Now the dataset only handles I/O (in thread workers), and all the transforms are applied to the batched tensor on the GPU in the main thread. Since tensorstore manages its own thread pool, it is recommended to set
num_workersto 1.