dynamically limiting the number of threads per proc

src/decentralizepy/mappings/Linear.py

@@ -59,3 +59,16 @@ class Linear(Mapping):
 
         """
         return (uid % self.procs_per_machine), (uid // self.procs_per_machine)
+
+    def get_local_procs_count(self):
+        """
+        Gives number of processes that run on the node
+
+        Returns
+        -------
+        int
+            the number of local processes
+
+        """
+
+        return self.procs_per_machine

src/decentralizepy/mappings/Mapping.py

@@ -67,3 +67,16 @@ class Mapping:
         """
 
         raise NotImplementedError
+
+    def get_local_procs_count(self):
+        """
+        Gives number of processes that run on the node
+
+        Returns
+        -------
+        int
+            the number of local processes
+
+        """
+
+        raise NotImplementedError

src/decentralizepy/node/Node.py

 import importlib
 import json
 import logging
+import math
 import os
 
 import torch
@@ -421,7 +422,9 @@ class Node:
             Other arguments
 
         """
-        torch.set_num_threads(2)
+        total_threads = os.cpu_count()
+        threads_per_proc = max(math.floor(total_threads / mapping.procs_per_machine), 1)
+        torch.set_num_threads(threads_per_proc)
         torch.set_num_interop_threads(1)
         self.instantiate(
             rank,
@@ -435,5 +438,8 @@ class Node:
             test_after,
             *args
         )
+        logging.info(
+            "Each proc uses %d threads out of %d.", threads_per_proc, total_threads
+        )
 
         self.run()