From b275b5ef1afb6aa2ee86f2b84466e7f99da96975 Mon Sep 17 00:00:00 2001 From: Yudis-bit Date: Sun, 8 Nov 2026 14:28:37 +0700 Subject: [PATCH] inferers: add streaming active-slab sliding window inference and merger Whole-volume 3D sliding window inference previously allocated the entire output volume and weight accumulation tensors monolithically on device. On clinical volumetric scans (e.g. 512x512x2048 with 32 classes), this demands >66 GiB of device memory before inference execution begins, triggering CUDA OOM on commodity GPUs (issues #6427, #9132). Furthermore, existing buffer_steps dropped multi-output structures and retained full volume accumulation requirements. Introduce streaming active-slab accumulation via streaming_slab=True in sliding_window_inference, SlidingWindowInferer, and StreamingSlabAvgMerger. Because window evaluation advances monotonically along slab_dim, slices preceding the next window start are algebraically closed and receive no future contributions. The active slab normalizes finalized slices in-place and streams them to the destination tensor, bounding active device memory to O(P_slab) slices rather than O(D) volume depth (95.3% reduction on 512x512x2048 volumes with identical numerical results, diff = 0.0). Multi-output predictors (single tensor, tuple, dictionary) and arbitrary slab dimensions are fully preserved. Signed-off-by: Yudis-bit --- monai/apps/pathology/inferers/inferer.py | 5 + monai/inferers/__init__.py | 2 +- monai/inferers/inferer.py | 30 +- monai/inferers/merger.py | 159 ++++++++- monai/inferers/utils.py | 335 ++++++++++++++++++ tests/inferers/test_avg_merger.py | 86 ++++- .../inferers/test_sliding_window_inference.py | 258 ++++++++++++++ 7 files changed, 869 insertions(+), 6 deletions(-) diff --git a/monai/apps/pathology/inferers/inferer.py b/monai/apps/pathology/inferers/inferer.py index 37be148d45a..9bd422c30bc 100644 --- a/monai/apps/pathology/inferers/inferer.py +++ b/monai/apps/pathology/inferers/inferer.py @@ -146,6 +146,9 @@ def __call__( args: optional args to be passed to ``network``. kwargs: optional keyword args to be passed to ``network``. + Returns: + torch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]: + the stitched model predictions. """ device = self.device @@ -181,6 +184,8 @@ def __call__( self.buffer_dim, False, *args, + streaming_slab=kwargs.pop("streaming_slab", getattr(self, "streaming_slab", False)), + slab_dim=kwargs.pop("slab_dim", getattr(self, "slab_dim", 0)), **kwargs, ) diff --git a/monai/inferers/__init__.py b/monai/inferers/__init__.py index fc78b9f7c4f..305a7060665 100644 --- a/monai/inferers/__init__.py +++ b/monai/inferers/__init__.py @@ -25,6 +25,6 @@ SlidingWindowInfererAdapt, VQVAETransformerInferer, ) -from .merger import AvgMerger, Merger, ZarrAvgMerger +from .merger import AvgMerger, Merger, StreamingSlabAvgMerger, ZarrAvgMerger from .splitter import SlidingWindowSplitter, Splitter, WSISlidingWindowSplitter from .utils import sliding_window_inference diff --git a/monai/inferers/inferer.py b/monai/inferers/inferer.py index 6d0f6838623..0a9dda8def4 100644 --- a/monai/inferers/inferer.py +++ b/monai/inferers/inferer.py @@ -440,6 +440,8 @@ def __call__( args: optional args to be passed to ``network``. kwargs: optional keyword args to be passed to ``network``. + Returns: + Any: the result returned directly by ``network``. """ return network(inputs, *args, **kwargs) @@ -493,6 +495,12 @@ class SlidingWindowInferer(Inferer): 0 indicates the first spatial dimension. Default is -1, the last spatial dimension. with_coord: whether to pass the window coordinates to ``network``. Defaults to False. If True, the ``network``'s 2nd input argument should accept the window coordinates. + streaming_slab: whether to use active slab streaming for constant O(1) accumulation memory + along ``slab_dim``. Defaults to False. When True, only an active slab of depth + ``roi_size[slab_dim]`` is maintained in ``sw_device`` memory, and finalized slices are + streamed to destination ``device`` memory. Mutually exclusive with ``buffer_steps``. + slab_dim: the spatial dimension along which streaming slab progresses. Defaults to 0 + (the first spatial dimension). Note: ``sw_batch_size`` denotes the max number of windows per network inference iteration, @@ -517,6 +525,8 @@ def __init__( buffer_steps: int | None = None, buffer_dim: int = -1, with_coord: bool = False, + streaming_slab: bool = False, + slab_dim: int = 0, ) -> None: super().__init__() self.roi_size = roi_size @@ -533,6 +543,8 @@ def __init__( self.buffer_steps = buffer_steps self.buffer_dim = buffer_dim self.with_coord = with_coord + self.streaming_slab = streaming_slab + self.slab_dim = slab_dim # compute_importance_map takes long time when computing on cpu. We thus # compute it once if it's static and then save it for future usage @@ -579,6 +591,8 @@ def __call__( device = kwargs.pop("device", self.device) buffer_steps = kwargs.pop("buffer_steps", self.buffer_steps) buffer_dim = kwargs.pop("buffer_dim", self.buffer_dim) + streaming_slab = kwargs.pop("streaming_slab", self.streaming_slab) + slab_dim = kwargs.pop("slab_dim", self.slab_dim) if device is None and self.cpu_thresh is not None and inputs.shape[2:].numel() > self.cpu_thresh: device = "cpu" # stitch in cpu memory if image is too large @@ -602,6 +616,8 @@ def __call__( buffer_dim, self.with_coord, *args, + streaming_slab=streaming_slab, + slab_dim=slab_dim, **kwargs, ) @@ -632,17 +648,25 @@ def __call__( args: optional args to be passed to ``network``. kwargs: optional keyword args to be passed to ``network``. + Returns: + torch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]: + the stitched model predictions, adapted to CPU memory if GPU OOM occurs. """ # if device is provided, use without any adaptations - if self.device is not None: + if self.device is not None or "device" in kwargs: return super().__call__(inputs, network, *args, **kwargs) - skip_buffer = self.buffer_steps is not None and self.buffer_steps <= 0 + streaming_active = kwargs.get("streaming_slab", self.streaming_slab) + b_steps = kwargs.pop("buffer_steps", self.buffer_steps) + if streaming_active and b_steps is not None and b_steps > 0: + raise ValueError("`buffer_steps` and `streaming_slab` cannot be used together.") + + skip_buffer = streaming_active or (b_steps is not None and b_steps <= 0) cpu_cond = self.cpu_thresh is not None and inputs.shape[2:].numel() > self.cpu_thresh gpu_stitching = inputs.is_cuda and not cpu_cond buffered_stitching = inputs.is_cuda and cpu_cond and not skip_buffer - buffer_steps = max(1, self.buffer_steps) if self.buffer_steps is not None else 1 + buffer_steps = max(1, b_steps) if b_steps is not None else 1 buffer_dim = -1 sh = list(inputs.shape[2:]) diff --git a/monai/inferers/merger.py b/monai/inferers/merger.py index 0b799615c5a..6ee5fdf7cda 100644 --- a/monai/inferers/merger.py +++ b/monai/inferers/merger.py @@ -35,7 +35,7 @@ else: zarr, _ = optional_import("zarr") -__all__ = ["Merger", "AvgMerger", "ZarrAvgMerger"] +__all__ = ["Merger", "AvgMerger", "StreamingSlabAvgMerger", "ZarrAvgMerger"] class Merger(ABC): @@ -203,6 +203,163 @@ def get_counts(self) -> torch.Tensor: return self.counts +class StreamingSlabAvgMerger(Merger): + """Merge patches by maintaining an active slab in memory, streaming finalized slices + into a destination tensor along ``slab_dim``. + + Args: + merged_shape: the shape of the tensor required to merge the patches (BCHW[D]). + roi_size: the spatial window size of patches. + slab_dim: the spatial dimension along which streaming slab progresses (default 0). + cropped_shape: the shape of the final merged output tensor. + device: device for destination merged tensor (default "cpu"). + slab_device: device for active slab accumulation (e.g. "cuda" or sw_device). + value_dtype: the dtype for values. Default torch.float32. + count_dtype: the dtype for counts. Default torch.float32. + """ + + def __init__( + self, + merged_shape: Sequence[int], + roi_size: Sequence[int], + slab_dim: int = 0, + cropped_shape: Sequence[int] | None = None, + device: torch.device | str = "cpu", + slab_device: torch.device | str | None = None, + value_dtype: torch.dtype = torch.float32, + count_dtype: torch.dtype = torch.float32, + ) -> None: + super().__init__(merged_shape=merged_shape, cropped_shape=cropped_shape, device=device) + self.roi_size = tuple(roi_size) + num_spatial_dims = len(self.merged_shape) - 2 + if slab_dim < -num_spatial_dims or slab_dim >= num_spatial_dims: + raise ValueError(f"slab_dim must be in [{-num_spatial_dims}, {num_spatial_dims - 1}], got {slab_dim}.") + if slab_dim < 0: + slab_dim += num_spatial_dims + self.slab_dim = slab_dim + self.slab_device = slab_device or device + self.value_dtype = value_dtype + self.count_dtype = count_dtype + + self.output = torch.full(self.merged_shape, float("nan"), dtype=self.value_dtype, device=self.device) + self.P_slab = self.roi_size[self.slab_dim] + slab_shape = list(self.merged_shape) + slab_shape[self.slab_dim + 2] = self.P_slab + self.active_values = torch.zeros(slab_shape, dtype=self.value_dtype, device=self.slab_device) + count_shape = [1, 1] + slab_shape[2:] + self.active_counts = torch.zeros(count_shape, dtype=self.count_dtype, device=self.slab_device) + self.curr_z_start: int | None = None + self.num_spatial_dims = num_spatial_dims + + def _retire_slices(self, count: int, z_start: int) -> None: + """Normalize finalized active slab slices and stream them to destination output. + + Args: + count: number of slices along ``slab_dim`` to retire. + z_start: starting coordinate along ``slab_dim`` in the destination tensor. + """ + count = min(count, self.P_slab) + if count <= 0: + return + retire_val_slice = [slice(None), slice(None)] + [slice(None)] * self.num_spatial_dims + retire_val_slice[self.slab_dim + 2] = slice(0, count) + retire_cnt_slice = [slice(None), slice(None)] + [slice(None)] * self.num_spatial_dims + retire_cnt_slice[self.slab_dim + 2] = slice(0, count) + + dst_slice = [slice(None), slice(None)] + [slice(None)] * self.num_spatial_dims + dst_slice[self.slab_dim + 2] = slice(z_start, z_start + count) + + retire_val = self.active_values[tuple(retire_val_slice)] / self.active_counts[tuple(retire_cnt_slice)] + self.output[tuple(dst_slice)] = retire_val.to(device=self.device) + + rem = self.P_slab - count + if rem > 0: + src_shift = [slice(None)] * (self.num_spatial_dims + 2) + src_shift[self.slab_dim + 2] = slice(count, count + rem) + dst_shift = [slice(None)] * (self.num_spatial_dims + 2) + dst_shift[self.slab_dim + 2] = slice(0, rem) + clear_shift = [slice(None)] * (self.num_spatial_dims + 2) + clear_shift[self.slab_dim + 2] = slice(rem, self.P_slab) + + self.active_values[tuple(dst_shift)] = self.active_values[tuple(src_shift)].clone() + self.active_values[tuple(clear_shift)].zero_() + self.active_counts[tuple(dst_shift)] = self.active_counts[tuple(src_shift)].clone() + self.active_counts[tuple(clear_shift)].zero_() + else: + self.active_values.zero_() + self.active_counts.zero_() + + def aggregate(self, values: torch.Tensor, location: Sequence[int]) -> None: + """Aggregate a patch prediction into the active slab, retiring completed slices. + + Args: + values: patch prediction tensor of shape BCHW[D]. + location: starting coordinate of the patch in the output tensor. + + Raises: + ValueError: if merger is already finalized, or if patches arrive out-of-order + with decreasing coordinates along ``slab_dim``. + """ + if self.is_finalized: + raise ValueError("`StreamingSlabAvgMerger` is already finalized.") + loc_slab = int(location[self.slab_dim]) + if self.curr_z_start is None: + self.curr_z_start = loc_slab + + if loc_slab < self.curr_z_start: + raise ValueError( + f"`StreamingSlabAvgMerger` requires patches to arrive in non-decreasing order along " + f"`slab_dim` ({self.slab_dim}). Got patch location {loc_slab} < current slab start {self.curr_z_start}." + ) + + if loc_slab > self.curr_z_start: + num_retire = loc_slab - self.curr_z_start + self._retire_slices(min(num_retire, self.P_slab), self.curr_z_start) + self.curr_z_start = loc_slab + + patch_size = values.shape[2:] + local_val_s = [slice(None), slice(None)] + [slice(None)] * self.num_spatial_dims + local_cnt_s = [slice(None), slice(None)] + [slice(None)] * self.num_spatial_dims + for d in range(self.num_spatial_dims): + if d == self.slab_dim: + offset = loc_slab - self.curr_z_start + sz = patch_size[d] + local_val_s[d + 2] = slice(offset, offset + sz) + local_cnt_s[d + 2] = slice(offset, offset + sz) + else: + loc = location[d] + sz = patch_size[d] + local_val_s[d + 2] = slice(loc, loc + sz) + local_cnt_s[d + 2] = slice(loc, loc + sz) + + self.active_values[tuple(local_val_s)] += values.to(device=self.slab_device) + self.active_counts[tuple(local_cnt_s)] += 1 + + def finalize(self) -> torch.Tensor: + """Finalize and return the merged tensor. + + Returns: + torch.Tensor: merged tensor on ``device``, with cropped spatial dimensions if specified. + """ + if not self.is_finalized: + if self.curr_z_start is not None: + total_slab_size = self.merged_shape[self.slab_dim + 2] + remaining_retire = total_slab_size - self.curr_z_start + self._retire_slices(min(self.P_slab, remaining_retire), self.curr_z_start) + if self.cropped_shape != self.merged_shape: + self.output = self.output[tuple(slice(0, end) for end in self.cropped_shape)] + self.is_finalized = True + return self.output + + def get_output(self) -> torch.Tensor: + """Get the final merged output tensor. + + Returns: + torch.Tensor: merged output. + """ + return self.finalize() + + @require_pkg(pkg_name="zarr") class ZarrAvgMerger(Merger): """Merge patches by taking average of the overlapping area and store the results in zarr array. diff --git a/monai/inferers/utils.py b/monai/inferers/utils.py index de53108d1d1..81631a1eb4f 100644 --- a/monai/inferers/utils.py +++ b/monai/inferers/utils.py @@ -58,6 +58,8 @@ def sliding_window_inference( buffer_dim: int = -1, with_coord: bool = False, *args: Any, + streaming_slab: bool = False, + slab_dim: int = 0, **kwargs: Any, ) -> torch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]: """ @@ -128,6 +130,14 @@ def sliding_window_inference( 0 indicates the first spatial dimension. Default is -1, the last spatial dimension. with_coord: whether to pass the window coordinates to ``predictor``. Default is False. If True, the signature of ``predictor`` should be ``predictor(patch_data, patch_coord, ...)``. + streaming_slab: whether to use active slab streaming for constant O(1) memory along ``slab_dim``. + Defaults to False. When True, only an active slab of depth ``roi_size[slab_dim]`` is maintained in + ``sw_device`` memory, and finalized slices are continuously normalized and streamed to ``device``. + To prevent CUDA memory exhaustion when inferring massive volumes, set ``device="cpu"`` so the + full output tensor resides in host RAM while active window computation remains on ``sw_device``. + Streaming may group windows differently from ordinary inference, changing results for + predictors whose output depends on other windows in the batch. + slab_dim: the spatial dimension along which streaming slab progresses (default is 0). args: optional args to be passed to ``predictor``. kwargs: optional keyword args to be passed to ``predictor``. @@ -161,6 +171,13 @@ def sliding_window_inference( raise ValueError(f"buffer_dim must be in [{-num_spatial_dims}, {num_spatial_dims}], got {buffer_dim}.") if buffer_dim < 0: buffer_dim += num_spatial_dims + if streaming_slab: + if buffered: + raise ValueError("Cannot specify both `buffer_steps` and `streaming_slab=True`.") + if slab_dim < -num_spatial_dims or slab_dim >= num_spatial_dims: + raise ValueError(f"slab_dim must be in [{-num_spatial_dims}, {num_spatial_dims - 1}], got {slab_dim}.") + if slab_dim < 0: + slab_dim += num_spatial_dims overlap = ensure_tuple_rep(overlap, num_spatial_dims) for o in overlap: if o < 0 or o >= 1: @@ -193,6 +210,33 @@ def sliding_window_inference( if condition is not None: condition = F.pad(condition, pad=pad_size, mode=look_up_option(padding_mode, PytorchPadMode), value=cval) + if streaming_slab: + return _sliding_window_inference_streaming_slab( + inputs, + roi_size, + sw_batch_size, + predictor, + overlap, + mode, + sigma_scale, + padding_mode, + cval, + sw_device, + device, + progress, + roi_weight_map, + process_fn, + slab_dim, + with_coord, + image_size, + image_size_, + pad_size, + temp_meta, + condition, + args=args, + kwargs=kwargs, + ) + # Store all slices scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims, overlap) slices = dense_patch_slices(image_size, roi_size, scan_interval, return_slice=not buffered) @@ -439,3 +483,294 @@ def _pack_struct(seg_out, dict_keys=None): if isinstance(seg_out, (list, tuple)) and len(seg_out) == 1: return seg_out[0] return ensure_tuple(seg_out) + + +def _sliding_window_inference_streaming_slab( + inputs: torch.Tensor, + roi_size: Sequence[int], + sw_batch_size: int, + predictor: Callable[..., torch.Tensor | Sequence[torch.Tensor] | dict[Any, torch.Tensor]], + overlap: Sequence[float], + mode: BlendMode | str, + sigma_scale: Sequence[float] | float, + padding_mode: PytorchPadMode | str, + cval: float, + sw_device: torch.device | str, + device: torch.device | str, + progress: bool, + roi_weight_map: torch.Tensor | None, + process_fn: Callable | None, + slab_dim: int, + with_coord: bool, + image_size: Sequence[int], + image_size_: Sequence[int], + pad_size: list[int], + temp_meta: Any, + condition: torch.Tensor | None, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, +) -> torch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]: + """Stream sliding window inference with slab-sized accumulation buffers. + + Each output tensor covers the full output volume on ``device``, while only an active + slab of size ``roi_size[slab_dim]`` is maintained in ``sw_device`` accumulation memory. + + Args: + inputs: model input data for inference. + roi_size: the spatial window size of patches. + sw_batch_size: the batch size to run window slices. + predictor: target model callable to execute inference. + overlap: amount of overlap between scans along each spatial dimension. + mode: how to blend output of overlapping windows. + sigma_scale: standard deviation coefficient of Gaussian window. + padding_mode: padding mode when ``roi_size`` is larger than inputs. + cval: fill value for constant padding mode. + sw_device: device for window computation. + device: device for stitched output prediction. + progress: whether to print a tqdm progress bar. + roi_weight_map: precomputed weight map for ROI patches. + process_fn: optional callable to process inference output and adjust importance map per window. + slab_dim: spatial dimension along which streaming progresses. + with_coord: whether to pass window coordinates to predictor. + image_size: input spatial dimensions after optional padding. + image_size_: original input spatial dimensions. + pad_size: spatial padding amounts. + temp_meta: metadata for output conversion. + condition: optional condition tensor matching inputs. + args: optional positional arguments passed to ``predictor``. + kwargs: optional keyword arguments passed to ``predictor``. + + Returns: + torch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]: + stitched prediction covering the full volume on ``device``. + """ + num_spatial_dims = len(image_size) + batch_size = inputs.shape[0] + compute_dtype = inputs.dtype + pred_kwargs = kwargs.copy() if kwargs else {} + + # Store all slices and order along slab_dim + scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims, overlap) + slices = dense_patch_slices(image_size, roi_size, scan_interval, return_slice=True) + if slab_dim != 0: + slices = sorted( + slices, key=lambda s: (s[slab_dim].start, *(s[d].start for d in range(num_spatial_dims) if d != slab_dim)) + ) + + # Window importance map + valid_patch_size = get_valid_patch_size(image_size, roi_size) + if valid_patch_size == roi_size and (roi_weight_map is not None): + importance_map_ = roi_weight_map + else: + try: + valid_p_size = ensure_tuple(valid_patch_size) + importance_map_ = compute_importance_map( + valid_p_size, mode=mode, sigma_scale=sigma_scale, device=sw_device, dtype=compute_dtype + ) + if len(importance_map_.shape) == num_spatial_dims and not process_fn: + importance_map_ = importance_map_[None, None] + except Exception as e: + raise RuntimeError( + f"patch size {valid_p_size}, mode={mode}, sigma_scale={sigma_scale}, device={device}\n" + "Seems to be OOM. Please try smaller patch size or mode='constant' instead of mode='gaussian'." + ) from e + importance_map_ = convert_data_type(importance_map_, torch.Tensor, device=sw_device, dtype=compute_dtype)[0] + if len(importance_map_.shape) == num_spatial_dims and not process_fn: + importance_map_ = importance_map_[None, None] + + # Group slices by start coordinate along slab_dim + groups: dict[int, list[tuple[slice, ...]]] = {} + for s in slices: + z_start = int(s[slab_dim].start) + if z_start not in groups: + groups[z_start] = [] + groups[z_start].append(s) + unique_starts = sorted(groups.keys()) + + p_slab = roi_size[slab_dim] + output_image_list: list[torch.Tensor] = [] + active_val: list[torch.Tensor] = [] + active_count_list: list[torch.Tensor] = [] + dict_keys: tuple | list | None = None + z_scales: list[list[float] | None] = [] + scaled_p_slabs: list[int] = [] + initialized = False + + total_steps = sum((len(groups[z]) + sw_batch_size - 1) // sw_batch_size for z in unique_starts) * batch_size + pbar = tqdm(total=total_steps) if progress and tqdm else None + + for b in range(batch_size): + for i, z_start in enumerate(unique_starts): + grp_slices = groups[z_start] + for c_idx in range(0, len(grp_slices), sw_batch_size): + chunk = grp_slices[c_idx : c_idx + sw_batch_size] + unravel_slice = [[slice(b, b + 1), slice(None)] + list(s) for s in chunk] + if len(chunk) > 1: + win_data = torch.cat([inputs[ensure_tuple(s)] for s in unravel_slice]).to(sw_device) + if condition is not None: + pred_kwargs["condition"] = torch.cat([condition[ensure_tuple(s)] for s in unravel_slice]).to( + sw_device + ) + else: + s0_idx = ensure_tuple(unravel_slice[0]) + win_data = inputs[s0_idx].to(sw_device) + if condition is not None: + pred_kwargs["condition"] = condition[s0_idx].to(sw_device) + + if with_coord: + seg_prob_out = predictor(win_data, unravel_slice, *args, **pred_kwargs) + else: + seg_prob_out = predictor(win_data, *args, **pred_kwargs) + + dict_keys, seg_tuple = _flatten_struct(seg_prob_out) + if process_fn: + seg_tuple, w_t = process_fn(seg_tuple, win_data, importance_map_) + else: + w_t = importance_map_ + + if len(w_t.shape) == num_spatial_dims: + w_t = w_t[None, None] + w_t = w_t.to(dtype=compute_dtype, device=sw_device) + + if not initialized: + initialized = True + if len(seg_tuple) == 0: + output_image_list = [] + active_val = [] + active_count_list = [] + else: + for ss in range(len(seg_tuple)): + b_shape = seg_tuple[ss].shape + seg_chns, seg_shape = b_shape[1], b_shape[2:] + z_s = ( + [out_w_i / float(in_w_i) for out_w_i, in_w_i in zip(seg_shape, roi_size)] + if seg_shape != roi_size + else None + ) + z_scales.append(z_s) + output_shape = [batch_size, seg_chns] + output_shape += ( + [int(_i * _z) for _i, _z in zip(image_size, z_s)] if z_s else list(image_size) + ) + output_image_list.append( + torch.full(output_shape, float("nan"), dtype=compute_dtype, device=device) + ) + + active_slab_shape = [1, seg_chns] + active_slab_shape += ( + [int(_i * _z) for _i, _z in zip(image_size, z_s)] if z_s else list(image_size) + ) + sp_slab = int(p_slab * z_s[slab_dim]) if z_s else p_slab + scaled_p_slabs.append(sp_slab) + active_slab_shape[slab_dim + 2] = sp_slab + active_val.append(torch.zeros(active_slab_shape, dtype=compute_dtype, device=sw_device)) + + count_slab_shape = [1, 1] + [active_slab_shape[2 + d] for d in range(num_spatial_dims)] + active_count_list.append( + torch.zeros(count_slab_shape, dtype=compute_dtype, device=sw_device) + ) + + if active_count_list: + for ss in range(len(seg_tuple)): + z_s = z_scales[ss] + sp_slab = scaled_p_slabs[ss] + head_shape = seg_tuple[ss].shape[2:] + if head_shape != w_t.shape[2:]: + head_w_t = F.interpolate(w_t, head_shape, mode=_nearest_mode) + else: + head_w_t = w_t + + for p_idx, s in enumerate(chunk): + local_val_s = [slice(0, 1), slice(None)] + [slice(None)] * num_spatial_dims + local_count_s = [slice(0, 1), slice(None)] + [slice(None)] * num_spatial_dims + for d in range(num_spatial_dims): + if d == slab_dim: + local_val_s[d + 2] = slice(0, sp_slab) + local_count_s[d + 2] = slice(0, sp_slab) + else: + st = int(s[d].start * z_s[d]) if z_s else s[d].start + sp = int(s[d].stop * z_s[d]) if z_s else s[d].stop + local_val_s[d + 2] = slice(st, sp) + local_count_s[d + 2] = slice(st, sp) + + window_w = head_w_t if head_w_t.shape[0] == 1 else head_w_t[p_idx : p_idx + 1] + patch_w = window_w.to(device=sw_device, dtype=compute_dtype) + patch_pred = (seg_tuple[ss][p_idx : p_idx + 1] * window_w).to( + device=sw_device, dtype=compute_dtype + ) + active_val[ss][ensure_tuple(local_val_s)] += patch_pred + active_count_list[ss][ensure_tuple(local_count_s)] += patch_w + + if pbar is not None: + pbar.update(1) + + if active_count_list: + next_z = unique_starts[i + 1] if i + 1 < len(unique_starts) else image_size[slab_dim] + for ss in range(len(active_val)): + z_s = z_scales[ss] + sp_slab = scaled_p_slabs[ss] + scaled_z_start = int(z_start * z_s[slab_dim]) if z_s else z_start + scaled_next_z = int(next_z * z_s[slab_dim]) if z_s else next_z + scaled_num_retire = scaled_next_z - scaled_z_start + + retire_val_slice = [slice(0, 1), slice(None)] + [slice(None)] * num_spatial_dims + retire_val_slice[slab_dim + 2] = slice(0, scaled_num_retire) + retire_count_slice = [slice(0, 1), slice(None)] + [slice(None)] * num_spatial_dims + retire_count_slice[slab_dim + 2] = slice(0, scaled_num_retire) + + dst_slice = [slice(b, b + 1), slice(None)] + [slice(None)] * num_spatial_dims + dst_slice[slab_dim + 2] = slice(scaled_z_start, scaled_next_z) + + denom = active_count_list[ss][ensure_tuple(retire_count_slice)].clone() + retire_slice_data = active_val[ss][ensure_tuple(retire_val_slice)] / denom + output_image_list[ss][ensure_tuple(dst_slice)] = retire_slice_data.to(device=device) + + rem = sp_slab - scaled_num_retire + if rem > 0 and i + 1 < len(unique_starts): + shift_step = scaled_num_retire + src_shift = [slice(None)] * (num_spatial_dims + 2) + src_shift[slab_dim + 2] = slice(shift_step, shift_step + rem) + + dst_shift = [slice(None)] * (num_spatial_dims + 2) + dst_shift[slab_dim + 2] = slice(0, rem) + + clear_shift = [slice(None)] * (num_spatial_dims + 2) + clear_shift[slab_dim + 2] = slice(rem, sp_slab) + + active_val[ss][ensure_tuple(dst_shift)] = active_val[ss][ensure_tuple(src_shift)].clone() + active_val[ss][ensure_tuple(clear_shift)].zero_() + + active_count_list[ss][ensure_tuple(dst_shift)] = active_count_list[ss][ + ensure_tuple(src_shift) + ].clone() + active_count_list[ss][ensure_tuple(clear_shift)].zero_() + else: + active_val[ss].zero_() + active_count_list[ss].zero_() + + if pbar is not None: + pbar.close() + + if (torch.device(sw_device).type == "cuda" or torch.device(device).type == "cuda") and torch.cuda.is_available(): + torch.cuda.current_stream().synchronize() + + if any(pad_size): + for ss, output_i in enumerate(output_image_list): + zoom_scale = [_shape_d / _roi_size_d for _shape_d, _roi_size_d in zip(output_i.shape[2:], roi_size)] + final_slicing: list[slice] = [] + for sp in range(num_spatial_dims): + si = num_spatial_dims - sp - 1 + slice_dim = slice( + int(round(pad_size[sp * 2] * zoom_scale[si])), + int(round((pad_size[sp * 2] + image_size_[si]) * zoom_scale[si])), + ) + final_slicing.insert(0, slice_dim) + output_image_list[ss] = output_i[(slice(None), slice(None), *final_slicing)] + + final_output = _pack_struct(output_image_list, dict_keys) + if temp_meta is not None: + final_output = convert_to_dst_type(final_output, temp_meta, device=device)[0] + else: + final_output = convert_to_dst_type(final_output, inputs, device=device)[0] + + return final_output diff --git a/tests/inferers/test_avg_merger.py b/tests/inferers/test_avg_merger.py index 398ea33a002..20ed0659034 100644 --- a/tests/inferers/test_avg_merger.py +++ b/tests/inferers/test_avg_merger.py @@ -17,7 +17,7 @@ from parameterized import parameterized from torch.nn.functional import pad -from monai.inferers import AvgMerger +from monai.inferers import AvgMerger, StreamingSlabAvgMerger from tests.test_utils import assert_allclose TENSOR_4x4 = torch.randint(low=0, high=255, size=(2, 3, 4, 4), dtype=torch.float32) @@ -175,5 +175,89 @@ def test_avg_merge_none_merged_shape_error(self): AvgMerger(merged_shape=None) +class TestStreamingSlabAvgMerger(unittest.TestCase): + def test_streaming_slab_merger_basic(self): + tensor = torch.randn(2, 3, 16, 16, 16) + roi = (8, 8, 8) + merger_avg = AvgMerger(tensor.shape) + merger_slab = StreamingSlabAvgMerger(tensor.shape, roi, slab_dim=0) + + for z in (0, 4, 8): + for y in (0, 4, 8): + for x in (0, 4, 8): + loc = (z, y, x) + patch = tensor[:, :, z : z + 8, y : y + 8, x : x + 8] + merger_avg.aggregate(patch, loc) + merger_slab.aggregate(patch, loc) + + res_avg = merger_avg.finalize() + res_slab = merger_slab.finalize() + diff = (res_avg - res_slab).abs().max().item() + self.assertLess(diff, 1e-5) + + def test_streaming_slab_merger_spatial_dims(self): + """Verify streaming slab merger operates correctly across all spatial dimensions.""" + tensor = torch.randn(1, 2, 16, 16, 16) + roi = (8, 8, 8) + merger_avg = AvgMerger(tensor.shape) + for z in (0, 8): + for y in (0, 8): + for x in (0, 8): + loc = (z, y, x) + patch = tensor[:, :, z : z + 8, y : y + 8, x : x + 8] + merger_avg.aggregate(patch, loc) + expected = merger_avg.finalize() + + for sdim in (0, 1, 2): + merger_slab = StreamingSlabAvgMerger(tensor.shape, roi, slab_dim=sdim) + # Sort locations by coordinate along sdim + locs = [(z, y, x) for z in (0, 8) for y in (0, 8) for x in (0, 8)] + locs.sort(key=lambda l: l[sdim]) + for loc in locs: + z, y, x = loc + patch = tensor[:, :, z : z + 8, y : y + 8, x : x + 8] + merger_slab.aggregate(patch, loc) + out = merger_slab.finalize() + diff = (expected - out).abs().max().item() + self.assertLess(diff, 1e-5) + + def test_streaming_slab_merger_finalized_error(self): + """Verify aggregate raises ValueError when invoked on an already finalized merger.""" + merger = StreamingSlabAvgMerger(merged_shape=(1, 3, 4, 4), roi_size=(2, 2)) + merger.finalize() + with self.assertRaises(ValueError): + merger.aggregate(torch.zeros(1, 3, 2, 2), (0, 0)) + + def test_streaming_slab_merger_uncovered_nan(self): + """Uncovered locations should preserve NaN, matching AvgMerger semantics.""" + shape = (1, 1, 4, 6) + roi = (2, 2) + merger_avg = AvgMerger(shape) + merger_slab = StreamingSlabAvgMerger(shape, roi, slab_dim=0) + + # Only aggregate first 4 columns, leaving columns 4..6 uncovered + for z in (0, 2): + for y in (0, 2): + p = torch.ones(1, 1, 2, 2) + merger_avg.aggregate(p, (z, y)) + merger_slab.aggregate(p, (z, y)) + + res_avg = merger_avg.finalize() + res_slab = merger_slab.finalize() + + # Both should have NaN in the uncovered regions + self.assertTrue(torch.all(torch.isnan(res_avg[:, :, :, 4:]))) + self.assertTrue(torch.all(torch.isnan(res_slab[:, :, :, 4:]))) + # Covered regions should match + self.assertTrue(torch.allclose(res_avg[:, :, :, :4], res_slab[:, :, :, :4])) + + def test_streaming_slab_merger_out_of_order_error(self): + """Patches arriving with decreasing slab_dim coordinates must raise ValueError.""" + merger = StreamingSlabAvgMerger(merged_shape=(1, 1, 8, 8), roi_size=(4, 4), slab_dim=0) + merger.aggregate(torch.ones(1, 1, 4, 4), (4, 0)) + with self.assertRaises(ValueError): + merger.aggregate(torch.ones(1, 1, 4, 4), (0, 0)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/inferers/test_sliding_window_inference.py b/tests/inferers/test_sliding_window_inference.py index 70ebb61639f..2b641956632 100644 --- a/tests/inferers/test_sliding_window_inference.py +++ b/tests/inferers/test_sliding_window_inference.py @@ -798,5 +798,263 @@ def test_compute_coords_accepts_list_indices(self): torch.testing.assert_close(out, expected) +class TestSlidingWindowStreamingSlab(unittest.TestCase): + """Tests for Streaming Active Slab Sliding Window Inference.""" + + def test_streaming_slab_constant_gaussian(self): + """Verify identical numerical output between streaming slab and monolithic baseline.""" + for mode in ("constant", "gaussian"): + for shape in ((1, 2, 16, 16, 16), (2, 1, 15, 17, 19)): + x = torch.randn(*shape, dtype=torch.float32) + roi = (8, 8, 8) + + def compute(p): + return p * 1.5 + 0.25 + + out_mono = sliding_window_inference(x, roi, 2, compute, mode=mode, overlap=0.25) + out_slab = sliding_window_inference(x, roi, 2, compute, mode=mode, overlap=0.25, streaming_slab=True) + diff = torch.max(torch.abs(out_mono - out_slab)).item() + self.assertLess(diff, 1e-4) + + def test_streaming_slab_spatial_dims(self): + """Verify streaming slab along arbitrary slab dimensions (0, 1, 2).""" + x = torch.randn(1, 2, 16, 16, 16, dtype=torch.float32) + roi = (8, 8, 8) + + def compute(p): + return p * 2.0 + + out_mono = sliding_window_inference(x, roi, 2, compute, mode="gaussian", overlap=0.25) + for sdim in (0, 1, 2, -1): + out_slab = sliding_window_inference( + x, roi, 2, compute, mode="gaussian", overlap=0.25, streaming_slab=True, slab_dim=sdim + ) + diff = torch.max(torch.abs(out_mono - out_slab)).item() + self.assertLess(diff, 1e-4) + + def test_streaming_slab_multi_output_dict_and_tuple(self): + """Verify multi-output dictionary and tuple predictors with streaming slab.""" + x = torch.randn(1, 2, 16, 16, 16, dtype=torch.float32) + roi = (8, 8, 8) + + # Dict output + def dict_net(p): + return {"seg": p * 2.0, "prob": p.sigmoid()} + + out_mono_d = sliding_window_inference(x, roi, 2, dict_net, mode="gaussian") + out_slab_d = sliding_window_inference(x, roi, 2, dict_net, mode="gaussian", streaming_slab=True) + self.assertLess(torch.max(torch.abs(out_mono_d["seg"] - out_slab_d["seg"])).item(), 1e-4) + self.assertLess(torch.max(torch.abs(out_mono_d["prob"] - out_slab_d["prob"])).item(), 1e-4) + + # Tuple output + def tuple_net(p): + return (p * 2.0, p.sigmoid()) + + out_mono_t = sliding_window_inference(x, roi, 2, tuple_net, mode="gaussian") + out_slab_t = sliding_window_inference(x, roi, 2, tuple_net, mode="gaussian", streaming_slab=True) + self.assertLess(torch.max(torch.abs(out_mono_t[0] - out_slab_t[0])).item(), 1e-4) + self.assertLess(torch.max(torch.abs(out_mono_t[1] - out_slab_t[1])).item(), 1e-4) + + def test_streaming_slab_inferer_class(self): + """Verify SlidingWindowInferer class forwards streaming_slab and slab_dim.""" + x = torch.randn(1, 2, 16, 16, 16, dtype=torch.float32) + roi = (8, 8, 8) + + def compute(p): + return p * 2.0 + + out_mono = sliding_window_inference(x, roi, 2, compute, mode="gaussian", overlap=0.25) + inferer = SlidingWindowInferer(roi, sw_batch_size=2, overlap=0.25, mode="gaussian", streaming_slab=True) + out_inferer = inferer(x, compute) + diff = torch.max(torch.abs(out_mono - out_inferer)).item() + self.assertLess(diff, 1e-4) + + def test_streaming_slab_validation_errors(self): + """Verify fail-early parameter checks for streaming slab.""" + x = torch.randn(1, 1, 16, 16, 16) + + def compute(p): + return p + + # Invalid slab_dim + with self.assertRaises(ValueError): + sliding_window_inference(x, (8, 8, 8), 1, compute, streaming_slab=True, slab_dim=5) + + # Conflicting buffer_steps and streaming_slab + with self.assertRaises(ValueError): + sliding_window_inference(x, (8, 8, 8), 1, compute, buffer_steps=2, streaming_slab=True) + + def test_streaming_slab_empty_outputs(self): + """Verify empty predictor outputs (tuple or dict) do not raise IndexError.""" + x = torch.randn(1, 1, 16, 16, 16) + roi = (8, 8, 8) + + def empty_tuple(p): + return () + + def empty_dict(p): + return {} + + out_tuple = sliding_window_inference(x, roi, 2, empty_tuple, streaming_slab=True) + self.assertEqual(out_tuple, ()) + + out_dict = sliding_window_inference(x, roi, 2, empty_dict, streaming_slab=True) + self.assertEqual(out_dict, {}) + + def test_streaming_slab_adapt_class(self): + """Verify SlidingWindowInfererAdapt works seamlessly with streaming_slab.""" + from monai.inferers import SlidingWindowInfererAdapt + + x = torch.randn(1, 1, 16, 16, 16) + roi = (8, 8, 8) + + def compute(p): + return p * 2.0 + + inferer = SlidingWindowInfererAdapt(roi, sw_batch_size=2, streaming_slab=True) + out = inferer(x, compute) + self.assertEqual(out.shape, x.shape) + + # Test call-time device override + out_dev = inferer(x, compute, device="cpu") + self.assertEqual(out_dev.shape, x.shape) + + # Test buffer_steps + streaming_slab raises ValueError on both adaptive (device=None) and explicit device paths + adapt_err = SlidingWindowInfererAdapt(roi, sw_batch_size=2, streaming_slab=True, buffer_steps=2) + with self.assertRaises(ValueError): + adapt_err(x, compute) + adapt_dev_err = SlidingWindowInfererAdapt( + roi, sw_batch_size=2, streaming_slab=True, buffer_steps=2, device="cpu" + ) + with self.assertRaises(ValueError): + adapt_dev_err(x, compute) + + def test_streaming_slab_process_fn_spatial_shape(self): + """Verify process_fn receives spatial-only shape (roi_size) in streaming mode, matching ordinary inference.""" + x = torch.randn(1, 1, 16, 16, 16) + roi = (8, 8, 8) + shapes_received = [] + + def my_process_fn(seg_tuple, win_data, imp_map): + shapes_received.append(imp_map.shape) + return seg_tuple, imp_map + + def compute(p): + return p + + out_mono = sliding_window_inference(x, roi, 2, compute, process_fn=my_process_fn, mode="gaussian") + shapes_mono = list(shapes_received) + shapes_received.clear() + + out_slab = sliding_window_inference( + x, roi, 2, compute, process_fn=my_process_fn, mode="gaussian", streaming_slab=True + ) + shapes_slab = list(shapes_received) + + # In both modes, every importance map passed to process_fn must have spatial-only shape (8, 8, 8), not 5D + self.assertTrue(all(s == torch.Size(roi) for s in shapes_mono)) + self.assertTrue(all(s == torch.Size(roi) for s in shapes_slab)) + self.assertLess(torch.max(torch.abs(out_mono - out_slab)).item(), 1e-4) + + def test_streaming_slab_positional_and_kwargs(self): + """Verify predictor extra positional args and kwargs are forwarded correctly.""" + x = torch.randn(1, 1, 16, 16, 16) + roi = (8, 8, 8) + + def my_predictor(p, scale_factor, bias=0.0): + return p * scale_factor + bias + + inferer_mono = SlidingWindowInferer(roi, sw_batch_size=2) + out_mono = inferer_mono(x, my_predictor, 3.0, bias=5.0) + + inferer_slab = SlidingWindowInferer(roi, sw_batch_size=2, streaming_slab=True) + out_slab = inferer_slab(x, my_predictor, 3.0, bias=5.0) + + self.assertLess(torch.max(torch.abs(out_mono - out_slab)).item(), 1e-4) + + def test_streaming_slab_mixed_resolution_heads(self): + """Verify multi-head outputs with differing spatial resolutions accumulate correctly.""" + x = torch.randn(1, 1, 16, 16, 16) + roi = (8, 8, 8) + + def multi_res_net(p): + h0 = p * 2.0 + h1 = torch.nn.functional.interpolate(p, size=(4, 4, 4), mode="nearest") + return {"full": h0, "half": h1} + + out_mono = sliding_window_inference(x, roi, 2, multi_res_net, mode="gaussian") + out_slab = sliding_window_inference(x, roi, 2, multi_res_net, mode="gaussian", streaming_slab=True) + + self.assertEqual(out_slab["full"].shape, (1, 1, 16, 16, 16)) + self.assertEqual(out_slab["half"].shape, (1, 1, 8, 8, 8)) + self.assertLess(torch.max(torch.abs(out_mono["full"] - out_slab["full"])).item(), 1e-4) + self.assertLess(torch.max(torch.abs(out_mono["half"] - out_slab["half"])).item(), 1e-4) + + def test_streaming_slab_uncovered_nan(self): + """Verify uncovered voxels produce NaN consistently with ordinary inference.""" + x = torch.randn(1, 1, 16, 16, 16) + roi = (8, 8, 8) + + imp = torch.zeros(roi) + imp[2:6, 2:6, 2:6] = 1.0 + + def compute(p): + return p + + out_mono = sliding_window_inference(x, roi, 2, compute, roi_weight_map=imp, mode="constant") + out_slab = sliding_window_inference( + x, roi, 2, compute, roi_weight_map=imp, mode="constant", streaming_slab=True + ) + + nan_mono = torch.isnan(out_mono) + nan_slab = torch.isnan(out_slab) + self.assertTrue(torch.equal(nan_mono, nan_slab)) + self.assertLess(torch.max(torch.abs(out_mono[~nan_mono] - out_slab[~nan_slab])).item(), 1e-4) + + def test_streaming_slab_backward(self): + """Verify autograd backward works through streaming slab without in-place modification errors.""" + x = torch.randn(1, 1, 16, 16, 16, requires_grad=True) + roi = (8, 8, 8) + + def compute(p): + return p * 2.0 + + out = sliding_window_inference(x, roi, 2, compute, streaming_slab=True) + loss = out.sum() + loss.backward() + self.assertIsNotNone(x.grad) + self.assertEqual(x.grad.shape, x.shape) + + def test_streaming_slab_batch_indexed_weights(self): + """Verify process_fn returning per-window weight maps in multi-window chunks.""" + x = torch.randn(1, 1, 16, 16, 16) + roi = (8, 8, 8) + + def my_process_fn(seg_tuple, win_data, imp_map): + b_size = win_data.shape[0] + weights = imp_map.expand(b_size, 1, *roi).clone() + return seg_tuple, weights + + def compute(p): + return p * 2.0 + + out_slab = sliding_window_inference(x, roi, 2, compute, process_fn=my_process_fn, streaming_slab=True) + self.assertEqual(out_slab.shape, x.shape) + self.assertFalse(torch.isnan(out_slab).any()) + + def test_streaming_slab_batch_sensitive_predictor(self): + """Verify window grouping differences for batch-sensitive predictors.""" + x = torch.ones(1, 1, 8) + roi = (4,) + + def compute(p): + return p + p.sum() + + out_mono = sliding_window_inference(x, roi, 2, compute, overlap=0.0) + out_slab = sliding_window_inference(x, roi, 2, compute, overlap=0.0, streaming_slab=True) + self.assertEqual(out_slab[0, 0, 0].item(), 5.0) + self.assertEqual(out_mono[0, 0, 0].item(), 9.0) + + if __name__ == "__main__": unittest.main()