From c49cb0a66085c5192245e9fc7e08ccc5bf74fe11 Mon Sep 17 00:00:00 2001 From: MikeUU332 Date: Thu, 18 Sep 2025 15:22:31 -0600 Subject: [PATCH] Update nondet_multi_threaded_augmenter.py Added global mutex instead of threadpool_limits --- .../dataloading/nondet_multi_threaded_augmenter.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/batchgenerators/dataloading/nondet_multi_threaded_augmenter.py b/batchgenerators/dataloading/nondet_multi_threaded_augmenter.py index 6a4819f..a8dd5a7 100755 --- a/batchgenerators/dataloading/nondet_multi_threaded_augmenter.py +++ b/batchgenerators/dataloading/nondet_multi_threaded_augmenter.py @@ -34,11 +34,14 @@ except ImportError: torch = None +from threading import Lock +global_mutex = Lock() def producer(queue: Queue, data_loader, transform, thread_id: int, seed, abort_event: Event, wait_time: float = 0.02): # the producer will set the abort event if something happens - with threadpool_limits(1, None): + # with threadpool_limits(1, None): + global_mutex.acquire() # using instead of threadpool_limits for thread safety np.random.seed(seed) data_loader.set_thread_id(thread_id) item = None @@ -72,7 +75,9 @@ def producer(queue: Queue, data_loader, transform, thread_id: int, seed, traceback.print_exc() abort_event.set() return - + + #End of mutex section + global_mutex.release() def pin_memory_of_all_eligible_items_in_dict(result_dict): for k in result_dict.keys(): @@ -276,4 +281,4 @@ def __del__(self): end = time() print(end - st) - mt._finish() \ No newline at end of file + mt._finish()