nessai.utils.threading

Utilities related to managing threads used by nessai.

Module Contents

Functions

configure_threads(max_threads=None, pytorch_threads=None, n_pool=None)

Configure the number of threads available.

Attributes

logger

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.