Sliding window inference for large 3D volumes? How to fit into memory?

[neuronflow]

Hi,

I want to run sliding window inference on big 3D microscopy images. My volumes have sizes of approximately (there is a bit of variance:

orig_size': [tensor([6656]), tensor([737]), tensor([5616])]

Obviously I cannot fit such a large image on the GPU at once, hence I thought using sliding window inference would be a good idea. My current try inspired by the 3D spleen tutorial:

    # inferer
    patch_size = crop_size
    sw_batch_size = 1
    overlap = 0.5

    inferer = SlidingWindowInferer(
        roi_size=patch_size,
        sw_batch_size=sw_batch_size,
        overlap=overlap,
        mode="gaussian",
        padding_mode="replicate",
   )


...
# loop through the inference data loader

inputs = data["images"].float()
outputs = inferer(inputs.to(device), model)

I set the batch_size to 1 for experiments. The network is trained on 96 x 96 x 64 crops, which I also use for the roi_size.

The code crashes with:

RuntimeError: CUDA out of memory. Tried to allocate 102.63 GiB (GPU 0; 47.46 GiB total capacity; 21.95 MiB already allocated; 46.11 GiB free; 26.00 MiB reserved in total by PyTorch)

From my naive understanding sliding window inference is exactly meant to fit such huge cases into memory? But I see that the code is expected to crash when the whole input tensor is sent to the CUDA device. How should I continue?

Obviously, I could try to slice my images manually into "digestible" patches with overlap and save the network outputs. And then reassemble the network outputs again based on my slicing grid. But such a sliding window inference of sliding window inferences seems counterintuitive to me?

Thanks for your help!

[wyli]

Yes, the SlidingWindowInferer provides options of sw_device and device. to address the memory issue I think you can do something like:

• set inputs.to('cpu')
• set model.to('cuda')
• then run the sliding window with SlidingWindowInferer(..., sw_device='cuda', device='cpu', ...)

so that only the model(window_data) part runs on cuda device, the rest is with cpu.

[phcerdan]

In this case, the outputs of the SlidingWindow will be in cpu or cuda? I am facing the following error when using the outputs in a DiceMetric.

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

[Nic-Ma]

Hi @phcerdan ,

The output device is consistent with device arg:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/inferers/inferer.py#L120

Thanks.

[neuronflow]

Thanks for the quick response. Okay, if I understand you correctly I should initialize like this?

    inferer = SlidingWindowInferer(
        roi_size=patch_size,
        sw_batch_size=sw_batch_size,
        sw_device="cuda",
        device="cpu",
        overlap=overlap,
        mode="gaussian",
        padding_mode="replicate",
    )

And remove the .to(device) part here?
outputs = inferer(inputs.to(device), model)

[wyli]

yes, and the model should be on "cuda" as well