AMASSS: Possible fix for CUDA memory issue (torch.cuda.OutOfMemoryError: CUDA out of memory)

[tschreiner]

Hi there,

as other people in this thread I am dealing with the CUDA running out of GPU memory (already reported in issue #17).

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.04 GiB (GPU 0; 8.00 GiB total capacity; 5.17 GiB already allocated; 673.43 MiB free; 5.30 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

The concrete exception happens when the function sliding_window_inference is called in AMASSS_CLI.py Line 976.

In a Project MONAI issue (https://github.com/Project-MONAI/MONAI/issues/1189) I found, that it is possible to use GPU computation but CPU memory by replacing

https://github.com/DCBIA-OrthoLab/SlicerAutomatedDentalTools/blob/86026595834472b132aca980b5cc0688359993c3/AMASSS_CLI/AMASSS_CLI.py#L976

by

val_outputs = sliding_window_inference(input_img, cropSize, args["nbr_GPU_worker"], net,overlap=args["precision"], sw_device="cuda", device="cpu")

in AMASS_CLI.py.

I have just tested it and now I am able to segment everything on a NVIDIA GeForce GTX 1060 6GB without any crashs due to no memory.

On a machine with Intel Core i7-6700 3.4 GHz CPU with 16GB memory and a NVIDIA GeForce GTX 1060 6GB a CBCT scan with a 0.4mm resolution took 780.25 seconds to segment (mandible, maxilla, cranial base and skin).

Can someone check and confirm the fix, please?

Thanks, Tedd