diff --git a/nnunetv2/run/run_training.py b/nnunetv2/run/run_training.py index 73a5ce88d..fff2ecd2e 100644 --- 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 index fc79558ba..c5b144a23 100644 --- 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 index 77c14d8d1..d61925466 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,8 @@ dependencies = [ "matplotlib", "seaborn", "imagecodecs", - "yacs" + "yacs", + "batchgeneratorsv2" ] [project.urls]