Sliding Window Inference with high gpu memory

[coreeey]

The size of mu input is 512x512x256, and i trained my 3D-UNet with rand crop 160x160x64，my aim is to segment 15 object ,so my outputs is 16x160x160x64，but i find that i can train my network with low gpu memory（4G）, but when infer with Sliding Window Inference , the gpu memory will rise to 24G and sometimes exceed the total memory, which really confused me a lot part of my code is below:
if cfg.val_amp is True:
with autocast():
val_outputs = sliding_window_inference(val_inputs, cfg.roi_size, cfg.sw_batch_size, model)
else:
val_outputs = sliding_window_inference(val_inputs, cfg.roi_size, cfg.sw_batch_size, model,)
where cfg.roi_size=(160,160,64) and cfg.sw_batch_size=1,
is there any solution to reduce the gpu memory?
thanks in advance

[wyli]

sure, the sliding window method takes sw_device and device arguments for the window data and the output volume respectively. for more info, please see https://docs.monai.io/en/stable/inferers.html#monai.inferers.sliding_window_inference

in this case you can set model and sw_device to a GPU, and device to a cpu

model = mode.to(device='cuda:0')
val_inputs = val_inputs.to(device="cpu")
val_outputs = sliding_window_inference(val_inputs, cfg.roi_size, cfg.sw_batch_size, model, device="cpu", sw_device="cuda:0")

[coreeey]

Thank u, i will take a try