Skip to content

Commit 229f519

Browse files
aymuos15ericspod
andauthored
Perf: skip redundant full-image mask on StdShiftIntensity nonzero=False path (#8975)
Fixes #8972 . ### Description On the default `nonzero=False` path `StdShiftIntensity._stdshift` built an all-True boolean mask the size of the image and shifted through it (`img[slices] = img[slices] + offset`), forcing a full advanced-index gather and scatter plus the mask allocation and `.any()` scan even though every voxel is selected. That is equivalent to shifting the image directly, so the `nonzero=False` branch now does `img + factor * std(img)`. Output is bit-for-bit identical and the `nonzero=True` path is untouched. `RandStdShiftIntensity` benefits too, since it also defaults to `nonzero=False`. Measured across 2D/3D, single and multi channel, float32 and float64 on both backends (best-of-3, CPU); output verified equal in every configuration: | shape | dtype | backend | current (ms) | proposed (ms) | speedup | |---|---|---|---|---|---| | 1x256x256 | f32 | numpy | 0.123 | 0.048 | 2.58x | | 1x256x256 | f32 | torch | 0.563 | 0.039 | 14.47x | | 1x1024x1024 | f32 | numpy | 3.016 | 1.294 | 2.33x | | 1x1024x1024 | f32 | torch | 7.040 | 0.274 | 25.66x | | 1x1024x1024 | f64 | torch | 9.644 | 0.601 | 16.04x | | 1x64x64x64 | f32 | numpy | 0.637 | 0.251 | 2.54x | | 1x64x64x64 | f64 | torch | 2.646 | 0.095 | 27.84x | | 1x128x128x128 | f32 | numpy | 7.019 | 2.914 | 2.41x | | 1x128x128x128 | f32 | torch | 41.618 | 0.630 | 66.09x | | 1x128x128x128 | f64 | torch | 44.534 | 1.825 | 24.40x | | 4x96x96x96 | f32 | torch | 68.152 | 1.357 | 50.23x | | 4x160x160x160 | f32 | numpy | 128.429 | 82.662 | 1.55x | | 4x160x160x160 | f32 | torch | 342.641 | 16.727 | 20.48x | | 4x160x160x160 | f64 | torch | 354.009 | 41.652 | 8.50x | numpy ranges 1.55x to 2.65x and torch 8.5x to 66x across the full sweep; the largest gains are on torch, where all-True boolean-mask indexing is especially costly. ### Types of changes - [x] Non-breaking change (fix or new feature that would not break existing functionality). --------- Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk> Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>
1 parent abbb47c commit 229f519

1 file changed

Lines changed: 2 additions & 5 deletions

File tree

monai/transforms/intensity/array.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -354,17 +354,14 @@ def __init__(
354354
self.dtype = dtype
355355

356356
def _stdshift(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
357-
ones: Callable
358357
std: Callable
359358
if isinstance(img, torch.Tensor):
360-
ones = torch.ones
361359
std = partial(torch.std, unbiased=False)
362360
else:
363-
ones = np.ones
364361
std = np.std
365362

366-
slices = (img != 0) if self.nonzero else ones(img.shape, dtype=bool)
367-
if slices.any():
363+
slices = (img != 0) if self.nonzero else ()
364+
if not self.nonzero or (isinstance(slices, (np.ndarray, torch.Tensor)) and slices.any()):
368365
offset = self.factor * std(img[slices])
369366
img[slices] = img[slices] + offset
370367
return img

0 commit comments

Comments
 (0)