wasserth / TotalSegmentator

Tool for robust segmentation of >100 important anatomical structures in CT and MR images
Apache License 2.0
1.48k stars 246 forks source link

Add option to force using CPU #37

Closed lassoan closed 1 year ago

lassoan commented 1 year ago

Background: TotalSegmentator works extremely well. It stands out among other segmentation tools with its robustness and the large number of segments, but also with its low resource needs - especially that it can run on a CPU. While all developers have access to computers with CUDA-capable GPU, most users (and basically anyone using a laptop) don't have CUDA-capable GPU. Uploading patient data to some remote server is not acceptable for many users. So, TotalSegmentator is awesome, it allows anyone to segment their own data locally, very conveniently within 3D Slicer, by a few clicks - and the Slicer community is really excited about this.

Problem: People who have a weak NVIDIA GPU (e.g., something like a 4GB GPU on a laptop) cannot currently get full-resolution segmentations because their GPU is detected (so CUDA is used) but they don't have enough GPU memory.

Would it be possible to add an option to TotalSegmentator to specify what device to use?

The options could be:

wasserth commented 1 year ago

Good idea. I will add this.

wasserth commented 1 year ago

Unfortunately it is a bit more complicated to do this since the relevant code parts are buried deep in the nnU-Net code. A very simple solution would be the following: If you want to run on CPU (even though you have a GPU) just install pytorch for CPU. Then TotalSegmentator will not find cuda inside of pytorch and therefore default to CPU. Can you define which pytorch to install in your 3d slicer extension?

lassoan commented 1 year ago

A very simple solution would be the following: If you want to run on CPU (even though you have a GPU) just install pytorch for CPU

This is what we currently recommend as a workaround. It is not ideal, because for smaller problems or fast mode, or for running other AI tools, the user may still want to take advantage of pytorch with GPU. Switching between CPU and CUDA pytorch is not seamless - the packages are big, installation takes significant time and not very robust, and may require restart of the application.

Can you define which pytorch to install in your 3d slicer extension?

We install pytorch using light-the-torch, which automatically finds the most suitable pytorch version.

Ideally, we would like to see a "device" option, similarly to how it is done for example in HD-BET.

Unfortunately it is a bit more complicated to do this since the relevant code parts are buried deep in the nnU-Net code.

I fully agree. This should be a feature of nnU-Net. Unfortunately, they don't seem to realize the impact of availability of GPU-free prediction. Requiring a CUDA-capable GPU essentially restricts the tool to run on developer computer or remote server, while if it can run on the CPU then it becomes universally available. The longer computation time does not matter much, because the one-minute fast mode is a good way for users to check if it is worth waiting for 30-60 minutes for a full-quality segmentation. If you submit a feature request to nnU-net and post the link here then I will help with making them understand the importance of this and hopefully they will work on it.

wasserth commented 1 year ago

I am thinking about using a environment variable to define the device and then inside of the nnUNet code at the different places this environment variable is checked. Would this work for you?

lassoan commented 1 year ago

An environment variable is a global variable within a process, which would force using the same device in a process. Therefore, it would be nicer to set the device as an argument for nnunet. However, the environment variable is not a big issue, as we currently start a separate process for each segmentation.

lassoan commented 1 year ago

Hundreds of people are now testing out TotalSegmentator in Slicer and about 5 reported this issue already. It would be nice to have a solution because manual reinstallation of pytorch is a much more complicated then the usual few-click setup.

wasserth commented 1 year ago

I tried to add the environment variable to nnU-Net to make it work with CPU, but unfortunately I am still getting errors. It seems that it is difficult in general to run a model completely on CPU, when torch is installed with CUDA. For example torch will use some CUDA implementation for the convolutional layers and when the input is a CPU tensor this will result in an error. I have the feeling you have to install torch for CPU to make it properly work on CPU.

lassoan commented 1 year ago

For example torch will use some CUDA implementation for the convolutional layers and when the input is a CPU tensor this will result in an error.

Most likely this is just a bug. With growing support for macOS MPL, Pytorch will surely clean up any hardcoded CUDA usages and allow properly specifying what devices should be used.

To confirm that this is indeed a bug, could you submit a bug report to pytorch and see how the developers respond? While you are preparing a minimal reproducible example for the bug report you may find out that the big is actually in nnunet, which would be very useful information, too (and make the problem easier to fix).

Thank you for all your work on this.

lipteck commented 1 year ago

Background: TotalSegmentator works extremely well. It stands out among other segmentation tools with its robustness and the large number of segments, but also with its low resource needs - especially that it can run on a CPU. While all developers have access to computers with CUDA-capable GPU, most users (and basically anyone using a laptop) don't have CUDA-capable GPU. Uploading patient data to some remote server is not acceptable for many users. So, TotalSegmentator is awesome, it allows anyone to segment their own data locally, very conveniently within 3D Slicer, by a few clicks - and the Slicer community is really excited about this.

Problem: People who have a weak NVIDIA GPU (e.g., something like a 4GB GPU on a laptop) cannot currently get full-resolution segmentations because their GPU is detected (so CUDA is used) but they don't have enough GPU memory.

Would it be possible to add an option to TotalSegmentator to specify what device to use?

The options could be:

  • automatic: the current method, use CUDA if available, otherwise CPU
  • CPU: force to use CPU, even if CUDA-capable GPU is available

Hi Wasserth

We have a HP z840 with GPU K40 10GB and K20 5GB, RAM 128 GB and 28 CPU on Ubuntu 20 LTS. We encountered error message saying these GPU are too old. Is there a quick workaround?

There is no result despite the fast option is selected in the slicer3d software.

Thanks!

regards Lip Teck

lassoan commented 1 year ago

If your GPU is not powerful enough then you may need to use the CPU-only version of pytorch as described here; https://github.com/lassoan/SlicerTotalSegmentator#setup

lipteck commented 1 year ago

Hi Lasso

I’ll try it. Meanwhile, I have a Dell 5820 with RTX A2000 on Windows 10. Do you think it is sufficient?

Thanks!

Regards Lip Teck

On Tue, 20 Dec 2022 at 9:47 PM, Andras Lasso @.***> wrote:

If your GPU is not powerful enough then you may need to use the CPU-only version of pytorch as described here; https://github.com/lassoan/SlicerTotalSegmentator#setup

— Reply to this email directly, view it on GitHub https://github.com/wasserth/TotalSegmentator/issues/37#issuecomment-1359380363, or unsubscribe https://github.com/notifications/unsubscribe-auth/AMZVTY2PCVBEUEK5PR72WVDWOG2FTANCNFSM6AAAAAAS7GMHR4 . You are receiving this because you commented.Message ID: @.***>

lassoan commented 1 year ago

How much GPU RAM do you have in your A2000?

lipteck commented 1 year ago

About 10-12 GB only. How much do we need?

On Tue, 20 Dec 2022 at 10:23 PM, Andras Lasso @.***> wrote:

How much GPU RAM do you have in your A2000?

— Reply to this email directly, view it on GitHub https://github.com/wasserth/TotalSegmentator/issues/37#issuecomment-1359441263, or unsubscribe https://github.com/notifications/unsubscribe-auth/AMZVTYZ6NQXSRGL56IWEUTLWOG6NTANCNFSM6AAAAAAS7GMHR4 . You are receiving this because you commented.Message ID: @.***>

lassoan commented 1 year ago

From nnU-Net documentation: "nnU-Net is guaranteed to run on GPUs with 11GB of memory. Many configurations may also run on 8 GB." - see more details here: https://github.com/MIC-DKFZ/nnUNet/blob/master/documentation/common_questions.md#can-i-run-nnu-net-on-smaller-gpus

lipteck commented 1 year ago

I have tried but problem persisted. Is there a command line to disable or disregard the GPU?

On Tue, 20 Dec 2022 at 9:47 PM, Andras Lasso @.***> wrote:

If your GPU is not powerful enough then you may need to use the CPU-only version of pytorch as described here; https://github.com/lassoan/SlicerTotalSegmentator#setup

— Reply to this email directly, view it on GitHub https://github.com/wasserth/TotalSegmentator/issues/37#issuecomment-1359380363, or unsubscribe https://github.com/notifications/unsubscribe-auth/AMZVTY2PCVBEUEK5PR72WVDWOG2FTANCNFSM6AAAAAAS7GMHR4 . You are receiving this because you commented.Message ID: @.***>

lassoan commented 1 year ago

To disregard the GPU, you need to install a CPU-only version of pytorch as described here: https://github.com/lassoan/SlicerTotalSegmentator#setup

I've updated the PyTorch util module (that takes care of installing PyTorch in Slicer) to make it easier to switch to CPU-only mode. You'll be able to select "cpu" as computation backend, and then uninstall/install PyTorch. It will be available in the Extensions Manager from tomorrow.

lipteck commented 1 year ago

Is this correct?

“/home/lipsteck/Slicer-5.2.1-linux-amd64/bin/PythonSlicer” -m pip install torch torchvision torch audio —force-reinstall

On Wed, 21 Dec 2022 at 9:39 PM, Andras Lasso @.***> wrote:

To disregard the GPU, you need to install a CPU-only version of pytorch as described here: https://github.com/lassoan/SlicerTotalSegmentator#setup

— Reply to this email directly, view it on GitHub https://github.com/wasserth/TotalSegmentator/issues/37#issuecomment-1361324581, or unsubscribe https://github.com/notifications/unsubscribe-auth/AMZVTYZ7LIFVSUXTLAJ2YL3WOMCAVANCNFSM6AAAAAAS7GMHR4 . You are receiving this because you commented.Message ID: @.***>

lassoan commented 1 year ago

You haven't used code formatting, therefore the text rendered here is incorrect, but your source text is probably OK. Just in case, I copy the text here with the correct formatting (using triple-backticks):

"/home/lipsteck/Slicer-5.2.1-linux-amd64/bin/PythonSlicer" -m pip install torch torchvision torchaudio --force-reinstall
lipteck commented 1 year ago

Do I need to use sudo?

"/home/lipteck/Slicer-5.2.1-linux-amd64/bin/PythonSlicer" -m pip install torch torchvision torch audio --force-reinstall

On Thu, 22 Dec 2022 at 10:11 AM, Andras Lasso @.***> wrote:

You haven't used code formatting, therefore the text rendered here is incorrect, but your source text is probably OK. Just in case, I copy the text here with the correct formatting (using triple-backticks):

"/home/lipsteck/Slicer-5.2.1-linux-amd64/bin/PythonSlicer" -m pip install torch torchvision torch audio --force-reinstall

— Reply to this email directly, view it on GitHub https://github.com/wasserth/TotalSegmentator/issues/37#issuecomment-1362307430, or unsubscribe https://github.com/notifications/unsubscribe-auth/AMZVTYZWJDJ65VCOPW5IAZLWOO2GPANCNFSM6AAAAAAS7GMHR4 . You are receiving this because you commented.Message ID: @.***>

lassoan commented 1 year ago

If your user has write access to the /home/lipteck/Slicer-5.2.1-linux-amd64 folder then you don't need sudo. If you have file access errors then probably Slicer is running - exit Slicer to make sure the files are not in use.

lipteck commented 1 year ago

Hi both

Just to update you that my attempt to install cpu version via the command line did not work. However, I have followed the instruction going to the PyTorch module to select cpu version to uninstall and install. It is now working!

Thank you!

I hope the code takes advantage of multiple cpu.

Happy New Year!

Regards Lip Teck

On Thu, 22 Dec 2022 at 11:28 AM, Andras Lasso @.***> wrote:

Closed #37 https://github.com/wasserth/TotalSegmentator/issues/37 as completed.

— Reply to this email directly, view it on GitHub https://github.com/wasserth/TotalSegmentator/issues/37#event-8092169953, or unsubscribe https://github.com/notifications/unsubscribe-auth/AMZVTY3PKHRIL4ZIQI5M4GDWOPDFJANCNFSM6AAAAAAS7GMHR4 . You are receiving this because you commented.Message ID: @.***>