Better GPU utilization during inference

[maunzzz]

Hi,
I'm using Monai's UNet for segmenting about 50 organs in CT images (~512x512x300).
I use sliding_window_inference to run the model on a new CT images, due to GPU memory limitations I need to keep the input and output volumes on cpu. The sliding window call (see below) is very slow and the GPU is poorly utilized. Using pytorch profiler and looking at the trace the GPU is active only 5% of the runtime.
Any tips and tricks to speed things up? The inference is done in mixed precision using pytorch's amp.

device = torch.device('gpu')
prediction = sliding_window_inference(inputs, # inputs is on cpu
                                      roi_size=(96, 96, 96),
                                      sw_batch_size=32,
                                      sw_device=device,
                                      predictor=self.model.network,
                                      overlap=0.25,
                                      device='cpu',
                                      mode='gaussian')

Thanks!

edit: formatting

[wyli]

there are some code optimisation recently, if you haven't tried the latest dev, please try installing from github (https://docs.monai.io/en/latest/installation.html#from-github), or use latest version of this file https://github.com/Project-MONAI/MONAI/blob/dev/monai/inferers/utils.py to replace the corresponding file in your installed version

[maunzzz]

Thank you for the answer!
Using the dev branch sliding_window_inference decreased the runtime from 16s to 14s for one image on my GeForce RTX 2070.
Still poor GPU utilization though, any other advice?

[wyli]

Could you share more details about pre/post processing if relevant? Cc @dongyang0122 @myron know more

[maunzzz]

I'm just timing the sliding_window_inference atm. The input image is float32.

[dongyang0122]

@maunzzz you can set the device in sliding_window_inference as gpu to generate the prediction on GPU devices instead of CPU. Communication between CPU and GPU may cost some time.

[maunzzz]

I have limited GPU memory and keeping the output on GPU causes OOM crash

[wyli]

we've been working on this, now with this implementation #6254 it's faster by default, and additional new options of buffer_steps and buffer_dim provide some flexibility in gpu output buffering. I used your configuration to benchmark and with a bit more memory usage the speed could be improved for a single volume:

#6254:

buffer step None buffer dim 1 MONAI SWI time: 4.219854116439819 sec, peak GPU mem:  0.955810546875 GB
buffer step 1 buffer dim 1 MONAI SWI time: 3.145899772644043 sec, peak GPU mem:  1.00341796875 GB
buffer step 3 buffer dim 1 MONAI SWI time: 2.9485719203948975 sec, peak GPU mem:  1.367919921875 GB
buffer step 5 buffer dim 1 MONAI SWI time: 2.8465116024017334 sec, peak GPU mem:  1.6126708984375 GB

monai v1.1:

MONAI SWI time: 5.842753648757935 sec, peak GPU mem:  0.94921875 GB

monai v1.2.0rc1:

MONAI SWI time: 4.93417763710022 sec, peak GPU mem:  0.94921875 GB

(these are mainly work from @myron and @dongyang0122)

[maunzzz]

Alright, that looks promising! I'll have a look at it. Thanks!

[wyli]

if along the buffer_dim there's no overlap between buffers, it'll enable the non-blocking mode and this may speed things up a bit, e.g when buffer_dim=-1, buffer_step=1:

overlap=(0.25,0.25,0), image size=(512,512,384), roi_size=96
buffer step 1 buffer dim -1 MONAI SWI time: 3.2060484886169434 sec, peak GPU mem:  1.08544921875 GB

overlap=(0.25,0.25,0), image_size=(512,512,385), roi_size=96
buffer step 1 buffer dim -1 MONAI SWI time: 4.261787414550781 sec, peak GPU mem:  1.08544921875 GB