Repository: Project-MONAI/MONAI

Issue:
About the sliding_window_inference
if the model has two or more outputs, the sliding_window_inference doesn't work,emm...


Discussion:
Hi @ZDYsszz ,

Thanks for your feedback.
Yes, the `sliding window` module doesn't support multiple outputs so far, it can support multiple inputs:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/inferers/inferer.py#L165

Thanks.


Patch:
diff --git a/monai/data/utils.py b/monai/data/utils.py
--- a/monai/data/utils.py
+++ b/monai/data/utils.py
@@ -974,7 +974,7 @@ def compute_importance_map(
     mode = look_up_option(mode, BlendMode)
     device = torch.device(device)
     if mode == BlendMode.CONSTANT:
-        importance_map = torch.ones(patch_size, device=device).float()
+        importance_map = torch.ones(patch_size, device=device, dtype=torch.float)
     elif mode == BlendMode.GAUSSIAN:
         center_coords = [i // 2 for i in patch_size]
         sigma_scale = ensure_tuple_rep(sigma_scale, len(patch_size))
@@ -987,15 +987,10 @@ def compute_importance_map(
         importance_map = importance_map.squeeze(0).squeeze(0)
         importance_map = importance_map / torch.max(importance_map)
         importance_map = importance_map.float()
-
-        # importance_map cannot be 0, otherwise we may end up with nans!
-        min_non_zero = importance_map[importance_map != 0].min().item()
-        importance_map = torch.clamp(importance_map, min=min_non_zero)
     else:
         raise ValueError(
             f"Unsupported mode: {mode}, available options are [{BlendMode.CONSTANT}, {BlendMode.CONSTANT}]."
         )
-
     return importance_map
 
 
diff --git a/monai/inferers/inferer.py b/monai/inferers/inferer.py
--- a/monai/inferers/inferer.py
+++ b/monai/inferers/inferer.py
@@ -9,13 +9,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Optional, Sequence, Union
+from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
 
 import torch
 import torch.nn as nn
 
-from monai.inferers.utils import sliding_window_inference
+from monai.inferers.utils import compute_importance_map, sliding_window_inference
 from monai.utils import BlendMode, PytorchPadMode, ensure_tuple
 from monai.visualize import CAM, GradCAM, GradCAMpp
 
@@ -120,6 +121,7 @@ class SlidingWindowInferer(Inferer):
             set to device=torch.device('cpu') the gpu memory consumption is less and independent of the
             `inputs` and `roi_size`. Output is on the `device`.
         progress: whether to print a tqdm progress bar.
+        cache_roi_weight_map: whether to precompute the ROI weight map.
 
     Note:
         ``sw_batch_size`` denotes the max number of windows per network inference iteration,
@@ -139,6 +141,7 @@ def __init__(
         sw_device: Union[torch.device, str, None] = None,
         device: Union[torch.device, str, None] = None,
         progress: bool = False,
+        cache_roi_weight_map: bool = False,
     ) -> None:
         Inferer.__init__(self)
         self.roi_size = roi_size
@@ -152,9 +155,26 @@ def __init__(
         self.device = device
         self.progress = progress
 
+        # 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
+        self.roi_weight_map = None
+        try:
+            if cache_roi_weight_map and isinstance(roi_size, Sequence) and min(roi_size) > 0:  # non-dynamic roi size
+                if device is None:
+                    device = "cpu"
+                self.roi_weight_map = compute_importance_map(
+                    ensure_tuple(self.roi_size), mode=mode, sigma_scale=sigma_scale, device=device
+                )
+            if cache_roi_weight_map and self.roi_weight_map is None:
+                warnings.warn("cache_roi_weight_map=True, but cache is not created. (dynamic roi_size?)")
+        except BaseException as e:
+            raise RuntimeError(
+                "Seems to be OOM. Please try smaller roi_size, or use mode='constant' instead of mode='gaussian'. "
+            ) from e
+
     def __call__(
         self, inputs: torch.Tensor, network: Callable[..., torch.Tensor], *args: Any, **kwargs: Any
-    ) -> torch.Tensor:
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor], Dict[Any, torch.Tensor]]:
         """
 
         Args:
@@ -165,7 +185,7 @@ def __call__(
             kwargs: optional keyword args to be passed to ``network``.
 
         """
-        return sliding_window_inference(
+        return sliding_window_inference(  # type: ignore
             inputs,
             self.roi_size,
             self.sw_batch_size,
@@ -178,6 +198,7 @@ def __call__(
             self.sw_device,
             self.device,
             self.progress,
+            self.roi_weight_map,
             *args,
             **kwargs,
         )
@@ -251,7 +272,7 @@ def __init__(self, spatial_dim: int = 0, *args, **kwargs) -> None:
 
     def __call__(
         self, inputs: torch.Tensor, network: Callable[..., torch.Tensor], *args: Any, **kwargs: Any
-    ) -> torch.Tensor:
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor], Dict[Any, torch.Tensor]]:
         """
         Args:
             inputs: 3D input for inference
diff --git a/monai/inferers/utils.py b/monai/inferers/utils.py
--- a/monai/inferers/utils.py
+++ b/monai/inferers/utils.py
@@ -9,13 +9,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Callable, List, Sequence, Tuple, Union
+import warnings
+from typing import Any, Callable, Dict, List, Mapping, Sequence, Tuple, Union
 
 import torch
 import torch.nn.functional as F
 
 from monai.data.utils import compute_importance_map, dense_patch_slices, get_valid_patch_size
-from monai.utils import BlendMode, PytorchPadMode, fall_back_tuple, look_up_option, optional_import
+from monai.transforms import Resize
+from monai.utils import (
+    BlendMode,
+    PytorchPadMode,
+    convert_data_type,
+    ensure_tuple,
+    fall_back_tuple,
+    look_up_option,
+    optional_import,
+)
 
 tqdm, _ = optional_import("tqdm", name="tqdm")
 
@@ -35,12 +45,21 @@ def sliding_window_inference(
     sw_device: Union[torch.device, str, None] = None,
     device: Union[torch.device, str, None] = None,
     progress: bool = False,
+    roi_weight_map: Union[torch.Tensor, None] = None,
     *args: Any,
     **kwargs: Any,
-) -> torch.Tensor:
+) -> Union[torch.Tensor, Tuple[torch.Tensor], Dict[Any, torch.Tensor]]:
     """
     Sliding window inference on `inputs` with `predictor`.
 
+    The outputs of `predictor` could be a tensor, a tuple, or a dictionary of tensors.
+    Each output in the tuple or dict value is allowed to have different resolutions with respect to the input.
+    e.g., the input patch spatial size is [128,128,128], the output (a tuple of two patches) patch sizes
+    could be ([128,64,256], [64,32,128]).
+    In this case, the parameter `overlap` and `roi_size` need to be carefully chosen to ensure the output ROI is still
+    an integer. If the predictor's input and output spatial sizes are not equal, we recommend choosing the parameters
+    so that `overlap*roi_size*output_size/input_size` is an integer (for each spatial dimension).
+
     When roi_size is larger than the inputs' spatial size, the input image are padded during inference.
     To maintain the same spatial sizes, the output image will be cropped to the original input size.
 
@@ -52,9 +71,16 @@ def sliding_window_inference(
             corresponding components of img size. For example, `roi_size=(32, -1)` will be adapted
             to `(32, 64)` if the second spatial dimension size of img is `64`.
         sw_batch_size: the batch size to run window slices.
-        predictor: given input tensor `patch_data` in shape NCHW[D], `predictor(patch_data)`
-            should return a prediction with the same spatial shape and batch_size, i.e. NMHW[D];
-            where HW[D] represents the patch spatial size, M is the number of output channels, N is `sw_batch_size`.
+        predictor: given input tensor ``patch_data`` in shape NCHW[D],
+            The outputs of the function call ``predictor(patch_data)`` should be a tensor, a tuple, or a dictionary
+            with Tensor values. Each output in the tuple or dict value should have the same batch_size, i.e. NM'H'W'[D'];
+            where H'W'[D'] represents the output patch's spatial size, M is the number of output channels,
+            N is `sw_batch_size`, e.g., the input shape is (7, 1, 128,128,128),
+            the output could be a tuple of two tensors, with shapes: ((7, 5, 128, 64, 256), (7, 4, 64, 32, 128)).
+            In this case, the parameter `overlap` and `roi_size` need to be carefully chosen
+            to ensure the scaled output ROI sizes are still integers.
+            If the `predictor`'s input and output spatial sizes are different,
+            we recommend choosing the parameters so that ``overlap*roi_size*zoom_scale`` is an integer for each dimension.
         overlap: Amount of overlap between scans.
         mode: {``"constant"``, ``"gaussian"``}
             How to blend output of overlapping windows. Defaults to ``"constant"``.
@@ -78,6 +104,8 @@ def sliding_window_inference(
             set to device=torch.device('cpu') the gpu memory consumption is less and independent of the
             `inputs` and `roi_size`. Output is on the `device`.
         progress: whether to print a `tqdm` progress bar.
+        roi_weight_map: pre-computed (non-negative) weight map for each ROI.
+            If not given, and ``mode`` is not `constant`, this map will be computed on the fly.
         args: optional args to be passed to ``predictor``.
         kwargs: optional keyword args to be passed to ``predictor``.
 
@@ -85,14 +113,14 @@ def sliding_window_inference(
         - input must be channel-first and have a batch dim, supports N-D sliding window.
 
     """
+    compute_dtype = inputs.dtype
     num_spatial_dims = len(inputs.shape) - 2
     if overlap < 0 or overlap >= 1:
-        raise AssertionError("overlap must be >= 0 and < 1.")
+        raise ValueError("overlap must be >= 0 and < 1.")
 
     # determine image spatial size and batch size
     # Note: all input images must have the same image size and batch size
-    image_size_ = list(inputs.shape[2:])
-    batch_size = inputs.shape[0]
+    batch_size, _, *image_size_ = inputs.shape
 
     if device is None:
         device = inputs.device
@@ -117,13 +145,27 @@ def sliding_window_inference(
     total_slices = num_win * batch_size  # total number of windows
 
     # Create window-level importance map
-    importance_map = compute_importance_map(
-        get_valid_patch_size(image_size, roi_size), mode=mode, sigma_scale=sigma_scale, device=device
-    )
+    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:
+            importance_map = compute_importance_map(valid_patch_size, mode=mode, sigma_scale=sigma_scale, device=device)
+        except BaseException as e:
+            raise RuntimeError(
+                "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=device, dtype=compute_dtype)[0]  # type: ignore
+    # handle non-positive weights
+    min_non_zero = max(importance_map[importance_map != 0].min().item(), 1e-3)
+    importance_map = torch.clamp(importance_map.to(torch.float32), min=min_non_zero).to(compute_dtype)
 
     # Perform predictions
-    output_image, count_map = torch.tensor(0.0, device=device), torch.tensor(0.0, device=device)
-    _initialized = False
+    dict_key, output_image_list, count_map_list = None, [], []
+    _initialized_ss = -1
+    is_tensor_output = True  # whether the predictor's output is a tensor (instead of dict/tuple)
+
+    # for each patch
     for slice_g in tqdm(range(0, total_slices, sw_batch_size)) if progress else range(0, total_slices, sw_batch_size):
         slice_range = range(slice_g, min(slice_g + sw_batch_size, total_slices))
         unravel_slice = [
@@ -131,31 +173,105 @@ def sliding_window_inference(
             for idx in slice_range
         ]
         window_data = torch.cat([inputs[win_slice] for win_slice in unravel_slice]).to(sw_device)
-        seg_prob = predictor(window_data, *args, **kwargs).to(device)  # batched patch segmentation
+        seg_prob_out = predictor(window_data, *args, **kwargs)  # batched patch segmentation
+
+        # convert seg_prob_out to tuple seg_prob_tuple, this does not allocate new memory.
+        if isinstance(seg_prob_out, torch.Tensor):
+            seg_prob_tuple = (seg_prob_out,)
+        elif isinstance(seg_prob_out, Mapping):
+            if dict_key is None:
+                dict_key = sorted(seg_prob_out.keys())  # track predictor's output keys
+            seg_prob_tuple = tuple(seg_prob_out[k] for k in dict_key)
+            is_tensor_output = False
+        else:
+            seg_prob_tuple = ensure_tuple(seg_prob_out)
+            is_tensor_output = False
+
+        # for each output in multi-output list
+        for ss, seg_prob in enumerate(seg_prob_tuple):
+            seg_prob = seg_prob.to(device)  # BxCxMxNxP or BxCxMxN
+
+            # compute zoom scale: out_roi_size/in_roi_size
+            zoom_scale = []
+            for axis, (img_s_i, out_w_i, in_w_i) in enumerate(
+                zip(image_size, seg_prob.shape[2:], window_data.shape[2:])
+            ):
+                _scale = out_w_i / float(in_w_i)
+                if not (img_s_i * _scale).is_integer():
+                    warnings.warn(
+                        f"For spatial axis: {axis}, output[{ss}] will have non-integer shape. Spatial "
+                        f"zoom_scale between output[{ss}] and input is {_scale}. Please pad inputs."
+                    )
+                zoom_scale.append(_scale)
 
-        if not _initialized:  # init. buffer at the first iteration
-            output_classes = seg_prob.shape[1]
-            output_shape = [batch_size, output_classes] + list(image_size)
-            # allocate memory to store the full output and the count for overlapping parts
-            output_image = torch.zeros(output_shape, dtype=torch.float32, device=device)
-            count_map = torch.zeros(output_shape, dtype=torch.float32, device=device)
-            _initialized = True
+            if _initialized_ss < ss:  # init. the ss-th buffer at the first iteration
+                # construct multi-resolution outputs
+                output_classes = seg_prob.shape[1]
+                output_shape = [batch_size, output_classes] + [
+                    int(image_size_d * zoom_scale_d) for image_size_d, zoom_scale_d in zip(image_size, zoom_scale)
+                ]
+                # allocate memory to store the full output and the count for overlapping parts
+                output_image_list.append(torch.zeros(output_shape, dtype=compute_dtype, device=device))
+                count_map_list.append(torch.zeros([1, 1] + output_shape[2:], dtype=compute_dtype, device=device))
+                _initialized_ss += 1
 
-        # store the result in the proper location of the full output. Apply weights from importance map.
-        for idx, original_idx in zip(slice_range, unravel_slice):
-            output_image[original_idx] += importance_map * seg_prob[idx - slice_g]
-            count_map[original_idx] += importance_map
+            # resizing the importance_map
+            resizer = Resize(spatial_size=seg_prob.shape[2:], mode="nearest", anti_aliasing=False)
+
+            # store the result in the proper location of the full output. Apply weights from importance map.
+            for idx, original_idx in zip(slice_range, unravel_slice):
+                # zoom roi
+                original_idx_zoom = list(original_idx)  # 4D for 2D image, 5D for 3D image
+                for axis in range(2, len(original_idx_zoom)):
+                    zoomed_start = original_idx[axis].start * zoom_scale[axis - 2]
+                    zoomed_end = original_idx[axis].stop * zoom_scale[axis - 2]
+                    if not zoomed_start.is_integer() or (not zoomed_end.is_integer()):
+                        warnings.warn(
+                            f"For axis-{axis-2} of output[{ss}], the output roi range is not int. "
+                            f"Input roi range is ({original_idx[axis].start}, {original_idx[axis].stop}). "
+                            f"Spatial zoom_scale between output[{ss}] and input is {zoom_scale[axis - 2]}. "
+                            f"Corresponding output roi range is ({zoomed_start}, {zoomed_end}).\n"
+                            f"Please change overlap ({overlap}) or roi_size ({roi_size[axis-2]}) for axis-{axis-2}. "
+                            "Tips: if overlap*roi_size*zoom_scale is an integer, it usually works."
+                        )
+                    original_idx_zoom[axis] = slice(int(zoomed_start), int(zoomed_end), None)
+                importance_map_zoom = resizer(importance_map.unsqueeze(0))[0].to(compute_dtype)
+                # store results and weights
+                output_image_list[ss][original_idx_zoom] += importance_map_zoom * seg_prob[idx - slice_g]
+                count_map_list[ss][original_idx_zoom] += (
+                    importance_map_zoom.unsqueeze(0).unsqueeze(0).expand(count_map_list[ss][original_idx_zoom].shape)
+                )
 
     # account for any overlapping sections
-    output_image = output_image / count_map
-
-    final_slicing: List[slice] = []
-    for sp in range(num_spatial_dims):
-        slice_dim = slice(pad_size[sp * 2], image_size_[num_spatial_dims - sp - 1] + pad_size[sp * 2])
-        final_slicing.insert(0, slice_dim)
-    while len(final_slicing) < len(output_image.shape):
-        final_slicing.insert(0, slice(None))
-    return output_image[final_slicing]
+    for ss in range(len(output_image_list)):
+        output_image_list[ss] = (output_image_list[ss] / count_map_list.pop(0)).to(compute_dtype)
+
+    # remove padding if image_size smaller than roi_size
+    for ss, output_i in enumerate(output_image_list):
+        if torch.isnan(output_i).any() or torch.isinf(output_i).any():
+            warnings.warn("Sliding window inference results contain NaN or Inf.")
+
+        zoom_scale = [
+            seg_prob_map_shape_d / roi_size_d for seg_prob_map_shape_d, roi_size_d in zip(output_i.shape[2:], roi_size)
+        ]
+
+        final_slicing: List[slice] = []
+        for sp in range(num_spatial_dims):
+            slice_dim = slice(pad_size[sp * 2], image_size_[num_spatial_dims - sp - 1] + pad_size[sp * 2])
+            slice_dim = slice(
+                int(round(slice_dim.start * zoom_scale[num_spatial_dims - sp - 1])),
+                int(round(slice_dim.stop * zoom_scale[num_spatial_dims - sp - 1])),
+            )
+            final_slicing.insert(0, slice_dim)
+        while len(final_slicing) < len(output_i.shape):
+            final_slicing.insert(0, slice(None))
+        output_image_list[ss] = output_i[final_slicing]
+
+    if dict_key is not None:  # if output of predictor is a dict
+        final_output = dict(zip(dict_key, output_image_list))
+    else:
+        final_output = tuple(output_image_list)
+    return final_output[0] if is_tensor_output else final_output  # type: ignore
 
 
 def _get_scan_interval(


Test patch:
diff --git a/tests/test_sliding_window_inference.py b/tests/test_sliding_window_inference.py
--- a/tests/test_sliding_window_inference.py
+++ b/tests/test_sliding_window_inference.py
@@ -177,6 +177,11 @@ def compute(self, data):
         )
         np.testing.assert_allclose(result.cpu().numpy(), expected, rtol=1e-4)
 
+        result = SlidingWindowInferer(
+            roi_shape, sw_batch_size, overlap=0.5, mode="gaussian", sigma_scale=[1.0, 1.0], cache_roi_weight_map=True
+        )(inputs, _Pred().compute)
+        np.testing.assert_allclose(result.cpu().numpy(), expected, rtol=1e-4)
+
     def test_cval(self):
         device = "cuda" if torch.cuda.is_available() else "cpu:0"
         inputs = torch.ones((1, 1, 3, 3)).to(device=device)
@@ -227,6 +232,7 @@ def compute(data, test1, test2):
             device,
             device,
             has_tqdm,
+            None,
             t1,
             test2=t2,
         )
@@ -238,6 +244,67 @@ def compute(data, test1, test2):
         )(inputs, compute, t1, test2=t2)
         np.testing.assert_allclose(result.cpu().numpy(), expected, rtol=1e-4)
 
+    def test_multioutput(self):
+        device = "cuda" if torch.cuda.is_available() else "cpu:0"
+        inputs = torch.ones((1, 6, 20, 20)).to(device=device)
+        roi_shape = (8, 8)
+        sw_batch_size = 10
+
+        def compute(data):
+            return data + 1, data[:, ::3, ::2, ::2] + 2, data[:, ::2, ::4, ::4] + 3
+
+        def compute_dict(data):
+            return {1: data + 1, 2: data[:, ::3, ::2, ::2] + 2, 3: data[:, ::2, ::4, ::4] + 3}
+
+        result = sliding_window_inference(
+            inputs,
+            roi_shape,
+            sw_batch_size,
+            compute,
+            0.5,
+            "constant",
+            1.0,
+            "constant",
+            0.0,
+            device,
+            device,
+            has_tqdm,
+            None,
+        )
+        result_dict = sliding_window_inference(
+            inputs,
+            roi_shape,
+            sw_batch_size,
+            compute_dict,
+            0.5,
+            "constant",
+            1.0,
+            "constant",
+            0.0,
+            device,
+            device,
+            has_tqdm,
+            None,
+        )
+        expected = (np.ones((1, 6, 20, 20)) + 1, np.ones((1, 2, 10, 10)) + 2, np.ones((1, 3, 5, 5)) + 3)
+        expected_dict = {1: np.ones((1, 6, 20, 20)) + 1, 2: np.ones((1, 2, 10, 10)) + 2, 3: np.ones((1, 3, 5, 5)) + 3}
+        for rr, ee in zip(result, expected):
+            np.testing.assert_allclose(rr.cpu().numpy(), ee, rtol=1e-4)
+        for rr, _ in zip(result_dict, expected_dict):
+            np.testing.assert_allclose(result_dict[rr].cpu().numpy(), expected_dict[rr], rtol=1e-4)
+
+        result = SlidingWindowInferer(
+            roi_shape, sw_batch_size, overlap=0.5, mode="constant", cval=-1, progress=has_tqdm
+        )(inputs, compute)
+        for rr, ee in zip(result, expected):
+            np.testing.assert_allclose(rr.cpu().numpy(), ee, rtol=1e-4)
+
+        result_dict = SlidingWindowInferer(
+            roi_shape, sw_batch_size, overlap=0.5, mode="constant", cval=-1, progress=has_tqdm
+        )(inputs, compute_dict)
+        for rr, _ in zip(result_dict, expected_dict):
+            np.testing.assert_allclose(result_dict[rr].cpu().numpy(), expected_dict[rr], rtol=1e-4)
+
 
 if __name__ == "__main__":
     unittest.main()