nessai.utils.threading
Utilities related to managing threads used by nessai.
Module Contents
Functions
|
Configure the number of threads available. |
Attributes
- nessai.utils.threading.logger
- nessai.utils.threading.configure_threads(max_threads=None, pytorch_threads=None, n_pool=None)
Configure the number of threads available.
This is necessary when using PyTorch on the CPU as by default it will use all available threads.
- Parameters
- max_threads: int, optional
Maximum total number of threads to use between PyTorch and multiprocessing.
- pytorch_threads: int, optional
Maximum number of threads for PyTorch on CPU.
- n_pool: int, optional
Number of pools to use if using multiprocessing.
Notes
Uses
torch.set_num_threads
. If pytorch threads is None but other arguments are specified then the value is inferred from them.