add control over the number of torch threads

pocomc/sampler.py

@@ -11,6 +11,7 @@
 from .flow import Flow
 from .particles import Particles
 from .geometry import Geometry
+from .threading import configure_threads
 
 class Sampler:
     r"""Preconditioned Monte Carlo class.
@@ -39,6 +40,9 @@ class Sampler:
     pool : pool
         Provided ``MPI`` or ``multiprocessing`` pool for
         parallelisation (default is ``pool=None``).
+    pytorch_threads : int
+        Maximum number of threads to use for torch. If ``None`` torch uses all
+        available threads while training the normalizing flow (default is ``pytorch_threads=1``).
     flow : ``torch.nn.Module`` or ``None``
         Normalizing flow (default is ``None``). The default is a Masked Autoregressive Flow
         (MAF) with 6 blocks of 3x64 layers and residual connections.
@@ -87,6 +91,7 @@ def __init__(self,
                  likelihood_kwargs: dict = None,
                  vectorize: bool = False,
                  pool=None,
+                 pytorch_threads=1,
                  flow=None,
                  train_config: dict = None,
                  precondition: bool = True,
@@ -105,6 +110,9 @@ def __init__(self,
             torch.manual_seed(random_state)
         self.random_state = random_state
 
+        # Configure threads
+        configure_threads(pytorch_threads=pytorch_threads)
+
         # Prior
         self.prior = prior
         self.log_prior = self.prior.logpdf

pocomc/threading.py

@@ -0,0 +1,21 @@
+import torch
+
+def configure_threads(pytorch_threads=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.
+
+    Notes
+    -----
+    Uses ``torch.set_num_threads``. If pytorch threads is None but other
+    arguments are specified then the value is inferred from them.
+
+    Parameters
+    ----------
+    pytorch_threads: int, optional
+        Maximum number of threads for PyTorch on CPU. If None, pytorch will
+        use all available threads.
+    """
+    if pytorch_threads:
+        torch.set_num_threads(pytorch_threads)