add environment variables for limiting number of threads

MIC-DKFZ · Apr 30, 2024 · 64df0a6 · 64df0a6
1 parent 85bf1f3
commit 64df0a6
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/nnunetv2/run/run_training.py b/nnunetv2/run/run_training.py
@@ -278,5 +278,8 @@ def run_training_entry():
 
 
 if __name__ == '__main__':
+    os.environ['OMP_NUM_THREADS'] = '1'
+    os.environ['MKL_NUM_THREADS'] = '1'
+    os.environ['OPENBLAS_NUM_THREADS'] = '1'
     # multiprocessing.set_start_method("spawn")
     run_training_entry()
diff --git a/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py b/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py
@@ -367,8 +367,6 @@ def _set_batch_size_and_oversample(self):
 
             print("worker", my_rank, "oversample", oversample_percent)
             print("worker", my_rank, "batch_size", batch_size_per_GPU[my_rank])
-            # self.print_to_log_file("worker", my_rank, "oversample", oversample_percents[my_rank])
-            # self.print_to_log_file("worker", my_rank, "batch_size", batch_sizes[my_rank])
 
             self.batch_size = batch_size_per_GPU[my_rank]
             self.oversample_foreground_percent = oversample_percent

diff --git a/pyproject.toml b/pyproject.toml
@@ -49,7 +49,8 @@ dependencies = [
     "matplotlib",
     "seaborn",
     "imagecodecs",
-    "yacs"
+    "yacs",
+    "batchgeneratorsv2"
 ]
 
 [project.urls]