Project-MONAI / Project-MONAI/MONAI
Memory efficient sliding window inference
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 8.7k
- Forks
- 1.6k
- Avg merge
- 5d 1h
- Merged PRs (30d)
- 20
Description
Is your feature request related to a problem? Please describe.
Large input volumes have to be processed via a sliding window algorithm, otherwise OOMs can happen quickly. There are two constraining properties which can cause an OOM: image size and number of predicted classes. The sliding_window_inference in MONAI allocates a FP32 probability aggregation buffer of size BxCxDxHxW and FP32 weight aggregation buffer of size BxDxHxW. In most cases, we are only interested in the final prediction anyway (class with highest probability).
The following example shows that even without running any model, the peak memory usage is very high (and that is a conservative image size if you think about e.g. high resolution isotropic wholebody CTs):
import torch
from monai.inferers import sliding_window_inference
num_classes = 100
data = torch.rand(1, 1, 384, 384, 256, dtype=torch.float32, device='cuda:0')
model = lambda x: x * torch.rand(x.shape[0], num_classes, x.shape[2], x.shape[3], x.shape[4], dtype=x.dtype, device=x.device)
out = sliding_window_inference(data, (128, 128, 128), 1, model)
|===========================================================================|
| PyTorch CUDA memory summary, device ID 0 |
|---------------------------------------------------------------------------|
| CUDA OOMs: 0 | cudaMalloc retries: 0 |
|===========================================================================|
| Metric | Cur Usage | Peak Usage | Tot Alloc | Tot Freed |
|---------------------------------------------------------------------------|
| Allocated memory | 14544 MB | 33504 MB | 166490 MB | 151946 MB |
| from large pool | 14544 MB | 33504 MB | 166490 MB | 151946 MB |
| from small pool | 0 MB | 0 MB | 0 MB | 0 MB |
|---------------------------------------------------------------------------|
| Active memory | 14544 MB | 33504 MB | 166490 MB | 151946 MB |
| from large pool | 14544 MB | 33504 MB | 166490 MB | 151946 MB |
| from small pool | 0 MB | 0 MB | 0 MB | 0 MB |
|---------------------------------------------------------------------------|
| GPU reserved memory | 35958 MB | 35958 MB | 35958 MB | 0 B |
| from large pool | 35956 MB | 35956 MB | 35956 MB | 0 B |
| from small pool | 2 MB | 2 MB | 2 MB | 0 B |
|---------------------------------------------------------------------------|
| Non-releasable memory | 0 B | 10853 MB | 12646 MB | 12646 MB |
| from large pool | 0 B | 10852 MB | 12638 MB | 12638 MB |
| from small pool | 0 B | 1 MB | 8 MB | 8 MB |
|---------------------------------------------------------------------------|
| Allocations | 2 | 10 | 221 | 219 |
| from large pool | 2 | 9 | 205 | 203 |
| from small pool | 0 | 3 | 16 | 16 |
|---------------------------------------------------------------------------|
| Active allocs | 2 | 10 | 221 | 219 |
| from large pool | 2 | 9 | 205 | 203 |
| from small pool | 0 | 3 | 16 | 16 |
|---------------------------------------------------------------------------|
| GPU reserved segments | 12 | 12 | 12 | 0 |
| from large pool | 11 | 11 | 11 | 0 |
| from small pool | 1 | 1 | 1 | 0 |
|---------------------------------------------------------------------------|
| Non-releasable allocs | 0 | 5 | 58 | 58 |
| from large pool | 0 | 4 | 53 | 53 |
| from small pool | 0 | 2 | 5 | 5 |
|---------------------------------------------------------------------------|
| Oversize allocations | 0 | 0 | 0 | 0 |
|---------------------------------------------------------------------------|
| Oversize GPU segments | 0 | 0 | 0 | 0 |
|===========================================================================|
If the image size or class count is further increased (e.g. 384x384x384) even on very powerful cards like A6000 with 48GB OOMs occur and require moving the aggregation to cpu memory. But since all computations for adding probability crops to the aggregation buffer are performed on the cpu instead of the gpu, it is going to be very slow.
Describe the solution you'd like
I implemented a sliding_window_inference_with_reduction method, which essentially performs two sliding windows. The outer loop iterates over a single dimension and performs the reduction operation, thus taking probabilities from the inner loop, apply e.g. argmax and store the results in the output buffer of size BxDxHxW with integer datatype (e.g. uint8). The inner loop performs a 2.5d sliding window inference of a slab of data. Since the outer loop can also iterate with some overlap, some of the probabilities of the inner loop will be used for initializing the buffer for the next inner loop iteration.
import torch
from inferer import sliding_window_inference_with_reduction
num_classes = 100
data = torch.rand(1, 1, 384, 384, 256, dtype=torch.float32, device='cuda:0')
model = lambda x: x * torch.rand(x.shape[0], num_classes, x.shape[2], x.shape[3], x.shape[4], dtype=x.dtype, device=x.device)
out = sliding_window_inference_with_reduction(data, (128, 128, 128), 1, model)
print(torch.cuda.memory_summary())
inferer.py
from typing import Any, Callable, Optional, Sequence, Tuple, Union
import torch
import torch.nn.functional as F
from monai.data.meta_tensor import MetaTensor
from monai.data.utils import (
compute_importance_map,
dense_patch_slices,
get_valid_patch_size,
)
from monai.inferers.utils import _get_scan_interval
from monai.utils import (
BlendMode,
PytorchPadMode,
convert_data_type,
convert_to_dst_type,
fall_back_tuple,
look_up_option,
)
from tqdm import tqdm
def sliding_window_inference_with_reduction(
inputs: torch.Tensor,
roi_size: Union[Sequence[int], int],
sw_batch_size: int,
predictor: Callable[..., torch.Tensor],
overlap: float = 0.25,
mode: Union[BlendMode, str] = BlendMode.CONSTANT,
sigma_scale: Union[Sequence[float], float] = 0.125,
padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
cval: float = 0.0,
sw_device: Optional[Union[torch.device, str]] = None,
device: Optional[Union[torch.device, str]] = None,
reduction_fn: Callable[..., torch.Tensor] = torch.argmax,
reduction_dim: int = 1,
output_dtype: torch.dtype = torch.uint8,
progress: bool = False,
*args: Any,
**kwargs: Any,
) -> torch.Tensor:
compute_dtype = inputs.dtype
num_spatial_dims = len(inputs.shape) - 2
if overlap < 0 or overlap >= 1:
raise ValueError("overlap must be >= 0 and < 1.")
batch_size, _, *orig_image_size = inputs.shape
if device is None:
device = inputs.device
if sw_device is None:
sw_device = inputs.device
roi_size_safe: Tuple[int] = fall_back_tuple(roi_size, orig_image_size)
image_size = tuple(
max(orig_image_size[i], roi_size_safe[i]) for i in range(num_spatial_dims)
)
pad_size = []
for k in range(len(inputs.shape) - 1, 1, -1):
diff = max(roi_size_safe[k - 2] - inputs.shape[k], 0)
half = diff // 2
pad_size.extend([half, diff - half])
if max(pad_size) > 0:
inputs = F.pad(
inputs,
pad=pad_size,
mode=look_up_option(padding_mode, PytorchPadMode),
value=cval,
)
patch_size = get_valid_patch_size(image_size, roi_size_safe)
importance_map = compute_importance_map(
patch_size=patch_size,
mode=mode,
sigma_scale=sigma_scale,
device=sw_device,
)
importance_map = torch.clamp(
importance_map,
min=max(importance_map[importance_map != 0].min().item(), 1e-3),
)
importance_map = convert_data_type(
importance_map, torch.Tensor, sw_device, compute_dtype
)[0]
# Identifiy outer and inner dimensions for the sliding window and aggregation
outer_dim = 2 # TODO Find heuristic for this
# Allocate buffers
output = torch.empty(
tuple(x for i, x in enumerate(inputs.shape) if i != reduction_dim),
dtype=output_dtype,
device=device,
)
slab_probabilities: Optional[torch.Tensor] = None
slab_weights = torch.zeros(
tuple(
inputs.shape[i]
if i != outer_dim + 2 # account for batch and channel dimensions
else roi_size_safe[outer_dim]
for i in range(inputs.ndim)
),
dtype=compute_dtype,
device=sw_device,
)
# Iterate over outer dimension and aggregate a full slab of reduced predictions
outer_step_size = int(roi_size_safe[outer_dim] * (1 - overlap))
outer_indices = list(
range(
0,
inputs.shape[outer_dim + 2] - roi_size_safe[outer_dim] + 1,
outer_step_size,
)
)
if outer_indices[-1] != image_size[outer_dim] - roi_size_safe[outer_dim]:
outer_indices.append(image_size[outer_dim] - roi_size_safe[outer_dim])
last_outer_dim_idx = -1
for outer_idx, outer_dim_idx in enumerate(
tqdm(outer_indices, leave=True, position=0) if progress else outer_indices
):
# Move old probabilities and weights based on the actual step size of this slab
if outer_idx > 0:
assert slab_probabilities is not None
actual_step_size = outer_dim_idx - last_outer_dim_idx
assert 0 < actual_step_size <= outer_step_size
new_slices = tuple(
slice(None)
if i != outer_dim + 2 # account only for batch dimension
else slice(None, -actual_step_size)
for i in range(slab_probabilities.ndim)
)
old_slices = tuple(
slice(None)
if i != outer_dim + 2 # account only for batch dimension
else slice(actual_step_size, None)
for i in range(slab_probabilities.ndim)
)
null_slices = tuple(
slice(None)
if i != outer_dim + 2 # account only for batch dimension
else slice(-actual_step_size, None)
for i in range(slab_probabilities.ndim)
)
slab_probabilities[new_slices] = slab_probabilities[old_slices]
slab_weights[new_slices] = slab_weights[old_slices]
slab_probabilities[null_slices] = 0.0
slab_weights[null_slices] = 0.0
# Take slab of input images and apply padding
slab_slices = tuple(
slice(None)
if i != outer_dim + 2 # account for batch and channel dimensions
else slice(outer_dim_idx, outer_dim_idx + roi_size_safe[outer_dim])
for i in range(inputs.ndim)
)
slab_input = inputs[slab_slices]
# Compute crop locations in slab
scan_interval = _get_scan_interval(
slab_input.shape[2:], roi_size_safe, num_spatial_dims, overlap
)
slices = dense_patch_slices(slab_input.shape[2:], roi_size_safe, scan_interval)
num_win = len(slices)
total_slices = num_win * batch_size
# Perform sliding window inference on slab
slice_indices = list(range(0, total_slices, sw_batch_size))
for slice_idx in (
tqdm(slice_indices, leave=False, position=1) if progress else slice_indices
):
# Get crops from slices
slice_range = range(slice_idx, min(slice_idx + sw_batch_size, total_slices))
unravel_slice = [
[slice(int(idx / num_win), int(idx / num_win) + 1), slice(None)]
+ list(slices[idx % num_win])
for idx in slice_range
]
window_data = torch.cat(
[
convert_data_type(slab_input[win_slice], torch.Tensor)[0]
for win_slice in unravel_slice
]
).to(sw_device)
# Compute probabilities and aggregate
probabilities = predictor(window_data, *args, **kwargs)
if slab_probabilities is None:
output_classes = probabilities.shape[1]
slab_probabilities = torch.zeros(
(batch_size, output_classes)
+ tuple(
image_size[i] if i != outer_dim else patch_size[outer_dim]
for i in range(len(image_size))
),
dtype=compute_dtype,
device=sw_device,
)
probabilities *= importance_map.unsqueeze(0).unsqueeze(0)
for slice_idx, win_slice in enumerate(unravel_slice):
slab_probabilities[win_slice] += probabilities[
slice_idx : slice_idx + 1
]
slab_weights[win_slice] += importance_map
# Apply reduction operation and move partial output to output buffer
assert slab_probabilities is not None
assert slab_weights is not None
copy_size = (
roi_size_safe[outer_dim]
if outer_idx == len(outer_indices) - 1
else outer_step_size
)
reduction_slices = tuple(
slice(None)
if i != outer_dim + 2 # account for batch and channel dimensions
else slice(None, copy_size)
for i in range(slab_probabilities.ndim)
)
predictions = reduction_fn(
slab_probabilities[reduction_slices] / slab_weights[reduction_slices],
dim=reduction_dim,
)
output[
tuple(
slice(None)
if i != outer_dim + 1 # account for batch dimension
else slice(outer_dim_idx, outer_dim_idx + copy_size)
for i in range(output.ndim)
)
] = predictions
last_outer_dim_idx = outer_dim_idx
# Crop to original image size
output = output[
..., : orig_image_size[0], : orig_image_size[1], : orig_image_size[2]
]
if isinstance(inputs, MetaTensor):
return convert_to_dst_type(output, inputs, device=device)[0]
return output
|===========================================================================|
| PyTorch CUDA memory summary, device ID 0 |
|---------------------------------------------------------------------------|
| CUDA OOMs: 0 | cudaMalloc retries: 0 |
|===========================================================================|
| Metric | Cur Usage | Peak Usage | Tot Alloc | Tot Freed |
|---------------------------------------------------------------------------|
| Allocated memory | 184320 KB | 16008 MB | 103790 MB | 103610 MB |
| from large pool | 184320 KB | 16008 MB | 103790 MB | 103610 MB |
| from small pool | 0 KB | 0 MB | 0 MB | 0 MB |
|---------------------------------------------------------------------------|
| Active memory | 184320 KB | 16008 MB | 103790 MB | 103610 MB |
| from large pool | 184320 KB | 16008 MB | 103790 MB | 103610 MB |
| from small pool | 0 KB | 0 MB | 0 MB | 0 MB |
|---------------------------------------------------------------------------|
| GPU reserved memory | 22486 MB | 22486 MB | 22486 MB | 0 B |
| from large pool | 22484 MB | 22484 MB | 22484 MB | 0 B |
| from small pool | 2 MB | 2 MB | 2 MB | 0 B |
|---------------------------------------------------------------------------|
| Non-releasable memory | 12288 KB | 5308 MB | 75658 MB | 75646 MB |
| from large pool | 12288 KB | 5308 MB | 75654 MB | 75642 MB |
| from small pool | 0 KB | 1 MB | 4 MB | 4 MB |
|---------------------------------------------------------------------------|
| Allocations | 2 | 11 | 168 | 166 |
| from large pool | 2 | 11 | 162 | 160 |
| from small pool | 0 | 3 | 6 | 6 |
|---------------------------------------------------------------------------|
| Active allocs | 2 | 11 | 168 | 166 |
| from large pool | 2 | 11 | 162 | 160 |
| from small pool | 0 | 3 | 6 | 6 |
|---------------------------------------------------------------------------|
| GPU reserved segments | 11 | 11 | 11 | 0 |
| from large pool | 10 | 10 | 10 | 0 |
| from small pool | 1 | 1 | 1 | 0 |
|---------------------------------------------------------------------------|
| Non-releasable allocs | 1 | 6 | 49 | 48 |
| from large pool | 1 | 6 | 46 | 45 |
| from small pool | 0 | 2 | 3 | 3 |
|---------------------------------------------------------------------------|
| Oversize allocations | 0 | 0 | 0 | 0 |
|---------------------------------------------------------------------------|
| Oversize GPU segments | 0 | 0 | 0 | 0 |
|===========================================================================|
In this example the peak memory usage is reduced in half. Further, it runs currently the outer loop on the third spatial dimension, so if you increase the size of this dimension the peak memory usage stays constant. E.g. 384x384x384 results in an OOM on A6000 with 48GB, but has still 16GB of peak memory usage using the optimized sliding window algorithm.
Current limitations before being able to integrate it into MONAI:
- Does not support multi-output models
- Does not support resizing outputs back to crop size
- Does not dynamically select the dimension for the outer loop (which could be an interesting heuristic for trade-off between GPU memory consumption and utilization)
Describe alternatives you've considered
Running sliding_window_inference with device='cpu', but it is of course significantly slower, since all computations regarding the aggregation buffer are performed on cpu instead of gpu.
Another approach could be to apply the same two step sliding window approach, but perform all the aggregation on the gpu and move the fully aggregated probabilities of a slab to the huge aggregation buffer in cpu memory. Then we have just a large copy operation instead of cpu-based computations for adding and normalizing probabilities.
Discussion
What do you think about this approach and would you like to integrate it into MONAI? In my opinion especially for inference or full image validation during training, where we only need the class with the highest probability, it would be a much more efficient implementation. This would allow remaining memory to be used for complex model computations instead of just a data storage. Maybe it could be further enhanced to also support applying multiple reductions, e.g. argmax and a FP measurement of uncertainty.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with MONAI's inferers.sliding_window_inference and the proposed inferer.py implementation in the issue. Review how the outer and inner sliding-window loops, reduction_fn, and output buffer are intended to interact, then determine the integration and tests needed for memory-efficient reduced output. Done means the feature handles the described large-volume and many-class case without the existing probability-buffer memory cost.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- computer-vision, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100