Repository: Project-MONAI/MONAI

Issue:
Memory error for sliding_window_inference
Hi MONAI-Team,
thanks for sharing your software development so far!
I tried to adapt the 3D spleen segmentation example to my own dataset. After some adjustments I got the training to run and as a next step tried to run some inference on my test data. The test volume is a 1504x1504x561 CT volume in nifti format.
To do so I used the sliding_window_inference on the CPU as the GPU immediatly runs out of memory.

```
test_transforms = Compose(
    [
        LoadNiftid(keys=["image", "label"]),
        AddChanneld(keys=["image", "label"]),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        CenterSpatialCropd(keys=["image", "label"], roi_size=[1000, 1000, 561]),
        ToTensord(keys=["image", "label"]),
    ]
)

test_ds = CacheDataset(data=test_files, transform=test_transforms, cache_rate=1, num_workers=1)
test_loader = DataLoader(test_ds, batch_size=1, num_workers=1)

test_device = torch.device("cpu")
model = UNet(
    dimensions=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(test_device)
loss_function = DiceLoss(to_onehot_y=True, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)

model.load_state_dict(torch.load(os.path.join(root_dir, "best_metric_model.pth")))
model.eval()

for test_data in test_loader:
    val_inputs, val_labels = (
        test_data["image"].to(test_device),
        test_data["label"].to(test_device),
    )
    print(f"batch_data image: {test_data['image'].shape}")
    print(f"batch_data label: {test_data['label'].shape}")
    roi_size = (64, 64, 64)
    sw_batch_size = 1
    test_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)
```

Unfortunately, I get either an out-of-memory error, a data loader error or the kernel dies. In each case the RAM of my machine (360 GB) runs out of memory during the sliding window inference. Therefore, I believe that all three errors are caused by the pilling up of the data in the RAM. 

What am I doing wrong? Did I miss interpreted something? I my understanding the ROI in the sliding window inferer should crop the sample to sub-volumes of the roi_size which then are used for inference by the network. Therefore, the memory footprint should be controlable by the roi_size and should even run on the GPU. But even if the entire sample is stored it should not require more than a couple of GBs (size of the nifti = 1.7 GB).  test_data does only consist of this single sample.
Thank you very much for your help!


Discussion:
thanks for the feedback, the sliding window memory size would be `1504x1504x561` x4 (4-byte float) x 2 (2channel prob) x 2 (tracking of spatial weights), so roughly 20G would be enough. 

but in your script, `CacheDataset` with cache rate =1  would take large memory and seems unnecessary, could you try using `Dataset`: something like `test_ds = monai.data.Dataset(data=test_files, transform=test_transforms)`
Hi wyli,
thank you very much for your quick response!
I did also try the normal dataset, unfortunately with the same result. Do you have an idea why the sliding_window_inference is building up so much more memory consumption (>300 GB) than anticipated? It takes up the space gradually while running the inference. 
Thank you very much!
sure, could you help run a few sanity checks:

for large tensors
```py
import torch
torch.rand((1, 5, 1504, 1504, 561))
```

for the sliding window
```py
from monai.inferers import sliding_window_inference
import torch
sliding_window_inference(torch.rand((1, 1, 1504, 1504, 561)), 64, 1, lambda x: torch.cat((x, x), axis=1)).shape
```
Sure, here is the ouput:
```
import torch
torch.rand((1, 5, 1504, 1504, 561))
```
tensor([[[[[3.3863e-01, 1.7919e-01, 6.0157e-01,  ..., 2.7129e-01,
            7.3396e-02, 4.2987e-02],
           [2.6321e-01, 9.2067e-01, 8.0036e-01,  ..., 4.5031e-01,
            4.1524e-01, 2.1147e-01],
           [5.1681e-01, 6.8210e-03, 2.0811e-01,  ..., 2.3153e-01,
            4.6916e-01, 6.0845e-01],
           ...,
           [1.4685e-02, 8.2935e-01, 6.7857e-01,  ..., 2.1001e-01,
            2.3881e-01, 2.0002e-01],
           [9.7772e-02, 3.7306e-01, 9.2843e-01,  ..., 3.5373e-01,
            7.8985e-01, 4.1494e-01],
           [5.9413e-01, 7.9738e-01, 4.6360e-01,  ..., 8.3947e-01,
            9.2945e-02, 7.0371e-01]]]]])

(obviously the tensor output is a lot longer but I copied the first and last lines)

```
from monai.inferers import sliding_window_inference
import torch
sliding_window_inference(torch.rand((1, 1, 1504, 1504, 561)), 64, 1, lambda x: torch.cat((x, x), axis=1)).shape
```
torch.Size([1, 2, 1504, 1504, 561])


Thank you again for your help!

thanks @joho84, so how much memory does it take to run the sliding window test? on my MacBook the peak memory was  ~20-30GB. I think the issue should be in the `test_ds` and `test_dataloader`. does it work if you comment out the last line of the script `test_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)`? I'll try to replicate the issue using some random images locally.
Hi wyli,
unfortunately I cannot check the memory consumption right now as I currently have no access to the machine. But when I ran the code yesterday it quickly terminated with no errors showing up. I will check the memory consumption as soon as I get a chance.
The dataloader itself works fine I think. I actually run the dataloader and the sliding_window_inference in different cells in the Jupyter Notebook. Running the first cells works fine and only consumes about 20 GB of memory:

```
test_transforms = Compose(
    [
        LoadNiftid(keys=["image", "label"]),
        AddChanneld(keys=["image", "label"]),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        CenterSpatialCropd(keys=["image", "label"], roi_size=[1000, 1000, 561]),
        ToTensord(keys=["image", "label"]),
    ]
)

test_ds = CacheDataset(data=test_files, transform=test_transforms, cache_rate=1, num_workers=1)
test_loader = DataLoader(test_ds, batch_size=1, num_workers=1)
```

Loading the network does also work:
```
test_device = torch.device("cpu")
model = UNet(
    dimensions=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(test_device)
loss_function = DiceLoss(to_onehot_y=True, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)
```

But running the sliding_window_inference leads to the explained behaviour:

```
model.load_state_dict(torch.load(os.path.join(root_dir, "best_metric_model.pth")))
model.eval()

for test_data in test_loader:
    val_inputs, val_labels = (
        test_data["image"].to(test_device),
        test_data["label"].to(test_device),
    )
    print(f"batch_data image: {test_data['image'].shape}")
    print(f"batch_data label: {test_data['label'].shape}")
    roi_size = (64, 64, 64)
    sw_batch_size = 1
    test_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)
```

Could you replicate the behaviour?

@wyli I just got a chance to  check the memory consumption:
for creating the random tensor: around 24 GB
for sliding_window_inference: peaking to around 34 GB but then dropping back to 0 GB
@wyli, @Nic-Ma: Thanks for the help so far! 
Could you replicate the memory error on your machine? Or do you have an idea what I should do differently?

Resolving patch:
diff --git a/monai/data/utils.py b/monai/data/utils.py
--- a/monai/data/utils.py
+++ b/monai/data/utils.py
@@ -482,7 +482,7 @@ def compute_importance_map(
     patch_size: Tuple[int, ...],
     mode: Union[BlendMode, str] = BlendMode.CONSTANT,
     sigma_scale: Union[Sequence[float], float] = 0.125,
-    device: Optional[torch.device] = None,
+    device: Union[torch.device, int, str] = "cpu",
 ) -> torch.Tensor:
     """Get importance map for different weight modes.
 
@@ -506,6 +506,7 @@ def compute_importance_map(
 
     """
     mode = BlendMode(mode)
+    device = torch.device(device)  # type: ignore[arg-type]
     if mode == BlendMode.CONSTANT:
         importance_map = torch.ones(patch_size, device=device).float()
     elif mode == BlendMode.GAUSSIAN:
diff --git a/monai/inferers/inferer.py b/monai/inferers/inferer.py
--- a/monai/inferers/inferer.py
+++ b/monai/inferers/inferer.py
@@ -87,7 +87,13 @@ class SlidingWindowInferer(Inferer):
             Padding mode when ``roi_size`` is larger than inputs. Defaults to ``"constant"``
             See also: https://pytorch.org/docs/stable/nn.functional.html#pad
         cval: fill value for 'constant' padding mode. Default: 0
-
+        sw_device: device for the window data.
+            By default the device (and accordingly the memory) of the `inputs` is used.
+            Normally `sw_device` should be consistent with the device where `predictor` is defined.
+        device: device for the stitched output prediction.
+            By default the device (and accordingly the memory) of the `inputs` is used. If for example
+            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`.
 
     Note:
         ``sw_batch_size`` denotes the max number of windows per network inference iteration,
@@ -104,6 +110,8 @@ def __init__(
         sigma_scale: Union[Sequence[float], float] = 0.125,
         padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
         cval: float = 0.0,
+        sw_device: Union[torch.device, str, None] = None,
+        device: Union[torch.device, str, None] = None,
     ) -> None:
         Inferer.__init__(self)
         self.roi_size = roi_size
@@ -113,6 +121,8 @@ def __init__(
         self.sigma_scale = sigma_scale
         self.padding_mode = padding_mode
         self.cval = cval
+        self.sw_device = sw_device
+        self.device = device
 
     def __call__(self, inputs: torch.Tensor, network: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor:
         """
@@ -133,4 +143,6 @@ def __call__(self, inputs: torch.Tensor, network: Callable[[torch.Tensor], torch
             sigma_scale=self.sigma_scale,
             padding_mode=self.padding_mode,
             cval=self.cval,
+            sw_device=self.sw_device,
+            device=self.device,
         )
diff --git a/monai/inferers/utils.py b/monai/inferers/utils.py
--- a/monai/inferers/utils.py
+++ b/monai/inferers/utils.py
@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, List, Optional, Sequence, Tuple, Union
+from typing import Callable, List, Sequence, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -28,7 +28,8 @@ def sliding_window_inference(
     sigma_scale: Union[Sequence[float], float] = 0.125,
     padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
     cval: float = 0.0,
-    device: Optional[torch.device] = None,
+    sw_device: Union[torch.device, str, None] = None,
+    device: Union[torch.device, str, None] = None,
 ) -> torch.Tensor:
     """
     Sliding window inference on `inputs` with `predictor`.
@@ -62,10 +63,13 @@ def sliding_window_inference(
             Padding mode for ``inputs``, when ``roi_size`` is larger than inputs. Defaults to ``"constant"``
             See also: https://pytorch.org/docs/stable/nn.functional.html#pad
         cval: fill value for 'constant' padding mode. Default: 0
-        device: device running the concatenation of the windows.
-            By default the device and accordingly the memory of the input device is used. If for example
+        sw_device: device for the window data.
+            By default the device (and accordingly the memory) of the `inputs` is used.
+            Normally `sw_device` should be consistent with the device where `predictor` is defined.
+        device: device for the stitched output prediction.
+            By default the device (and accordingly the memory) of the `inputs` is used. If for example
             set to device=torch.device('cpu') the gpu memory consumption is less and independent of the
-            input and roi_size parameter. Output is on the device set or if not set the inputs device.
+            `inputs` and `roi_size`. Output is on the `device`.
 
     Note:
         - input must be channel-first and have a batch dim, supports N-D sliding window.
@@ -81,6 +85,8 @@ def sliding_window_inference(
 
     if device is None:
         device = inputs.device
+    if sw_device is None:
+        sw_device = inputs.device
 
     roi_size = fall_back_tuple(roi_size, image_size_)
     # in case that image size is smaller than roi size
@@ -113,7 +119,7 @@ def sliding_window_inference(
             [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([inputs[win_slice] for win_slice in unravel_slice])
+        window_data = torch.cat([inputs[win_slice] for win_slice in unravel_slice]).to(sw_device)
         seg_prob = predictor(window_data).to(device)  # batched patch segmentation
 
         if not _initialized:  # init. buffer at the first iteration


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
@@ -16,6 +16,7 @@
 from parameterized import parameterized
 
 from monai.inferers import SlidingWindowInferer, sliding_window_inference
+from tests.utils import skip_if_no_cuda
 
 TEST_CASES = [
     [(2, 3, 16), (4,), 3, 0.25, "constant", torch.device("cpu:0")],  # 1D small roi
@@ -85,6 +86,21 @@ def compute(data):
         expected_val = np.ones((1, 3, 16, 15, 7), dtype=np.float32) + 1
         np.testing.assert_allclose(result.cpu().numpy(), expected_val)
 
+    @skip_if_no_cuda
+    def test_sw_device(self):
+        inputs = torch.ones((1, 3, 16, 15, 7)).to(device="cpu")
+        roi_shape = (4, 10, 7)
+        sw_batch_size = 10
+
+        def compute(data):
+            self.assertEqual(data.device.type, "cuda")
+            return data + torch.tensor(1, device="cuda")
+
+        result = sliding_window_inference(inputs, roi_shape, sw_batch_size, compute, sw_device="cuda")
+        np.testing.assert_string_equal(inputs.device.type, result.device.type)
+        expected_val = np.ones((1, 3, 16, 15, 7), dtype=np.float32) + 1
+        np.testing.assert_allclose(result.cpu().numpy(), expected_val)
+
     def test_sigma(self):
         device = "cuda" if torch.cuda.is_available() else "cpu:0"
         inputs = torch.ones((1, 1, 7, 7)).to(device=device)