tbepler / topaz

Pipeline for particle picking in cryo-electron microscopy images using convolutional neural networks trained from positive and unlabeled examples. Also featuring micrograph and tomogram denoising with DNNs.
GNU General Public License v3.0
172 stars 64 forks source link

Number of threads #51

Closed biochem-fan closed 4 years ago

biochem-fan commented 4 years ago

Even with topaz preprocess -t 1, the process uses all cores. One has to use OMP_NUM_THREADS to control the actual number of threads. It would be nice if one could specify the number of threads and number of worker processes.

The tutorial 01 is not correct:

-t/--num-workers X sets preprocess to use X threads, this and GPU device are mutually exclusive

tbepler commented 4 years ago

The number of CPU threads used is controlled by pytorch. Last I checked the only reliable way to set this was using OMP_NUM_THREADS. The --num-workers argument controls how many python processes are used, so how many micrographs are preprocessed in parallel.

The description in the tutorial is fixed in commit ba31c208800394af364a47af4092ddc69f118c6e.

biochem-fan commented 4 years ago

Great! Did you try setting OMP_NUM_THREADS within Python, before importing pytorch?

tbepler commented 4 years ago

Update: I just added an argument for setting the number of threads to the relevant commands (train, segment, extract, preprocess, normalize, denoise) in commit 9ee040e5adef8f907f160af07895ffdf84e07b2e. This sets the number of threads through a pytorch function. The caveat is that this function is broken in a few older pytorch versions, which is why I recommend using OMP_NUM_THREADS as the most bullet proof way for setting the number of threads. It works for me with pytorch version 1.1 and should work in the most recent versions of pytorch, but it's something to keep in mind.

Closing this issue since this commit adds the desired functionality.