set openmp threads to half the logical cores

NOEL-MNI · May 3, 2022 · cdc2701 · cdc2701
1 parent 394310f
commit cdc2701
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/app/inference.py b/app/inference.py
@@ -18,10 +18,14 @@
 if options['cuda'].startswith('cuda1'):
     os.environ["THEANO_FLAGS"] = "mode=FAST_RUN,device=cuda1,floatX=float32,dnn.enabled=False"
 elif options['cuda'].startswith('cpu'):
-    os.environ['OMP_NUM_THREADS'] = str(multiprocessing.cpu_count())
-    os.environ['MKL_NUM_THREADS'] = str(multiprocessing.cpu_count())
-    os.environ['GOTO_NUM_THREADS'] = str(multiprocessing.cpu_count())
-    os.environ['openmp'] = 'True'
+    os.environ['OMP_NUM_THREADS'] = str(multiprocessing.cpu_count() // 2)
+    var = os.getenv('OMP_NUM_THREADS', None)
+    try:
+        print("# of threads initialized: {}".format(int(var)))
+    except ValueError:
+        raise TypeError("The environment variable OMP_NUM_THREADS"
+                        " should be a number, got '%s'." % var)
+    # os.environ['openmp'] = 'True'
     os.environ["THEANO_FLAGS"] = "mode=FAST_RUN,device=cpu,openmp=True,floatX=float32"
 else:
     os.environ["THEANO_FLAGS"] = "mode=FAST_RUN,device=cuda0,floatX=float32,dnn.enabled=False"