Limit Number of threads in python SimpleITK programatically

I’m using SimpleITK from python. By default all logical CPUs which are found are used, but I like to constrict this. I already read this thread here: Threading in ITK and can confirm that setting the environment variable ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS does work.
However, I would like to be able to set this from the program itself, so I don’t have to mess with the variables all the time.

I found out it is possible to set this inside the python program like:

import os
import psutil
os.environ['ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS'] = str(psutil.cpu_count(logical=False))
import SimpleITK as sitk

But this has to be before importing SimpleITK.
Is it possible to set this also after SimpleITK was imported? I tried importlib.reload(sitk), but this seems to change nothing, the old number of Threads is still used.

Hello @reox,

The SimpleITK process object has a method SetGlobalDefaultNumberOfThreads:

import SimpleITK as sitk
MAX_THREADS = 3

sitk.ProcessObject.SetGlobalDefaultNumberOfThreads(MAX_THREADS)

There’s the matching Get method so that you can find out what is the global default number of threads.

3 Likes

@zivy
great, that works! thank you!

1 Like

Sorry for vamping this old thread, but it might be relevant.

@zivy, do you recommend setting this to 1 in a multiprocessing context? (e.g. PyTorch DataLoader with num_workers > 1).

1 Like

Hello @fepegar,

No worries about reviving an old thread. I do recommend setting this to 1 when doing multiprocessing.

I’ve had bad experiences when multi-threading and multi-processing are done at the same time. It was painful enough that I documented it here. I now realize we should make this much more visible and will add it to the SimpleITK FAQ, so thanks for the question.

3 Likes

The other thing to remember with multi-processing in python is that is can be implemented in multiple ways. If the multi-processing is done with a “fork” from the main process the ITK number of threads will be set in the subprocess. However, with other schedulers like some available for Dask, a whole new process can be initialized with and consideration for how the new process is setup and important need to be made.

Thank you, @zivy. I’ve been suspecting for a while that this is the case, so I’m glad to hear this clear guidance.

Thanks, @blowekamp. I’m not sure I follow 100%. For someone that doesn’t know much about concurrency, do you have suggestions on how/whether to set the number of threads for each scheduler/multiprocessing context?