issues with sliding window inference

[jxchen01]

I am trying to do inference of a 3D UNet model on a very large image. Because the image is large, I need to use sliding window inference. But, I am not sure how to best use this function. Basically, even reducing the sliding window patch size to a very small one, I still get out-of-memory issue. (I am using A100 GPU with 40GB memory)

Here is my code to reproduce this issue. Is there anything I am doing wrong?

import numpy as np
import torch
from random import random
from monai.networks.nets import UNet
from monai.inferers import sliding_window_inference

tmm = torch.cuda.get_device_properties(0).total_memory
rmm = torch.cuda.memory_reserved(0)
amm = torch.cuda.memory_allocated(0)
print(f"at the beginning, we have {rmm} reserved memory, {amm} allocated memory, out of {tmm/1000000000} GB")

model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=[16, 32, 64, 128, 256],
    strides=[2, 2, 2, 2],
    num_res_units=2
)

model.eval().cuda()

tmm = torch.cuda.get_device_properties(0).total_memory
rmm = torch.cuda.memory_reserved(0)
amm = torch.cuda.memory_allocated(0)
print(f"After creating the model, we have {rmm/1000000000} GB reserved memory, {amm/1000000000} GB allocated memory, out of {tmm/1000000000}GB")


# make some fake data
x = np.random.rand(1, 1, 400, 2048, 2048)

with torch.no_grad():
    # TODO: add convert to tensor with proper type, similar to torchio check
    x = torch.tensor(x).float().cuda()

    tmm = torch.cuda.get_device_properties(0).total_memory
    rmm = torch.cuda.memory_reserved(0)
    amm = torch.cuda.memory_allocated(0)
    print(f"After loading the image to GPU and before running sliding window, we have {rmm/1000000000} GB reserved memory, {amm/1000000000} GB allocated memory, out of {tmm/1000000000} GB")

    y_hat = sliding_window_inference(
        inputs=x,
        predictor=model,
        roi_size=[64, 64, 64],
        sw_batch_size=1,
        overlap=0.1,
        mode="gaussian",
    )

What I got was as follows:

at the beginning, we have 0 reserved memory, 0 allocated memory, out of 42.50632192 GB
After creating the model, we have 0.027262976 GB reserved memory, 0.019241984 GB allocated memory, out of 42.50632192GB
After loading the image to GPU and before running sliding window, we have 6.738149376 GB reserved memory, 6.730128384 GB allocated memory, out of 42.50632192 GB

Traceback (most recent call last):
  File "/mnt/eternus/users/Jianxu/projects/peyers_patch/test_OOM.py", line 41, in <module>
    y_hat = sliding_window_inference(
  File "/mnt/data/ISAS.DE/jianxu.chen/anaconda3/envs/paper/lib/python3.9/site-packages/monai/inferers/utils.py", line 248, in sliding_window_inference
    output_image_list[ss] = (output_image_list[ss] / count_map_list.pop(0)).to(compute_dtype)
RuntimeError: CUDA out of memory. Tried to allocate 12.50 GiB (GPU 0; 39.59 GiB total capacity; 31.27 GiB already allocated; 6.73 GiB free; 31.28 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

Any suggestion on any workaround?

Many thanks!
Jianxu

[jxchen01]

Okay .... Just find the solution, setting device=torch.device('cpu') will do the job.