From 11d56f8e844de629adde7c70caf5116294848481 Mon Sep 17 00:00:00 2001 From: MikeUU332 Date: Thu, 18 Sep 2025 15:16:16 -0600 Subject: [PATCH 1/2] Update data_loader.py Added global mutex fix for background workers issue. --- nnunetv2/training/dataloading/data_loader.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/nnunetv2/training/dataloading/data_loader.py b/nnunetv2/training/dataloading/data_loader.py index 0c51b8272..4095ebd35 100644 --- a/nnunetv2/training/dataloading/data_loader.py +++ b/nnunetv2/training/dataloading/data_loader.py @@ -15,6 +15,8 @@ from nnunetv2.utilities.plans_handling.plans_handler import PlansManager from acvl_utils.cropping_and_padding.bounding_boxes import crop_and_pad_nd +from threading import Lock +global_mutex = Lock() class nnUNetDataLoader(DataLoader): def __init__(self, @@ -198,7 +200,8 @@ def generate_train_batch(self): if self.transforms is not None: with torch.no_grad(): - with threadpool_limits(limits=1, user_api=None): + # with threadpool_limits(limits=1, user_api=None): + global_mutex.acquire() data_all = torch.from_numpy(data_all).float() seg_all = torch.from_numpy(seg_all).to(torch.int16) images = [] @@ -213,6 +216,7 @@ def generate_train_batch(self): else: seg_all = torch.stack(segs) del segs, images + global_mutex.release() return {'data': data_all, 'target': seg_all, 'keys': selected_keys} return {'data': data_all, 'target': seg_all, 'keys': selected_keys} From 9298023ab97eef6c70fdbdffbf0b12c3fd9145e5 Mon Sep 17 00:00:00 2001 From: MikeUU332 Date: Fri, 19 Sep 2025 16:38:38 -0600 Subject: [PATCH 2/2] Update data_loader.py formatting --- nnunetv2/training/dataloading/data_loader.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nnunetv2/training/dataloading/data_loader.py b/nnunetv2/training/dataloading/data_loader.py index 4095ebd35..2e214233d 100644 --- a/nnunetv2/training/dataloading/data_loader.py +++ b/nnunetv2/training/dataloading/data_loader.py @@ -217,6 +217,7 @@ def generate_train_batch(self): seg_all = torch.stack(segs) del segs, images global_mutex.release() + return {'data': data_all, 'target': seg_all, 'keys': selected_keys} return {'data': data_all, 'target': seg_all, 'keys': selected_keys}