Setup multiprocessing environment variables

quadra/configs/hydra/default.yaml

@@ -12,3 +12,9 @@ hydra:
         exclude_keys:
           - trainer.devices
     chdir: true
+    env_set:
+      OMP_NUM_THREADS: 8
+      OPENBLAS_NUM_THREADS: 8
+      MKL_NUM_THREADS: 8
+      VECLIB_MAXIMUM_THREADS: 8
+      NUMEXPR_NUM_THREADS: 8

quadra/main.py

@@ -2,30 +2,35 @@
 
 import hydra
 from omegaconf import DictConfig
-from pytorch_lightning import seed_everything
 
-from quadra.tasks.base import Task
 from quadra.utils.resolver import register_resolvers
-from quadra.utils.utils import get_logger, load_envs, setup_opencv
-from quadra.utils.validator import validate_config
 
-load_envs()
 register_resolvers()
 
-log = get_logger(__name__)
-
 
+# pylint: disable=import-outside-toplevel
 @hydra.main(config_path="configs/", config_name="config.yaml", version_base="1.3.0")
 def main(config: DictConfig):
     """Main entry function for any of the tasks."""
+    # Import everything here to allow setting environment variables before importing torch
+    import torch
+    from pytorch_lightning import seed_everything
+
+    from quadra.tasks import Task
+    from quadra.utils import utils
+    from quadra.utils.utils import get_logger, load_envs, setup_opencv
+    from quadra.utils.validator import validate_config
+
+    log = get_logger(__name__)
+    log.debug("Running with %d torch threads: ", torch.get_num_threads())
+
+    load_envs()
     if config.validate:
         start = time.time()
         validate_config(config)
         stop = time.time()
         log.info("Config validation took %f seconds", stop - start)
 
-    from quadra.utils import utils  # pylint: disable=import-outside-toplevel
-
     utils.extras(config)
 
     # Prints the resolved configuration to the console