Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optional Torch Multiprocessing in nnUNet for Improved Security and Compatibility #2614

Open
wants to merge 6 commits into
base: master
Choose a base branch
from
Open

Optional Torch Multiprocessing in nnUNet for Improved Security and Compatibility #2614

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
69fd9a7
Non-mp predict support [WIP]
surajpaib Nov 21, 2024
a595f64
add torchification
LennyN95 Nov 21, 2024
6cb4e7a
update preprocess_fromfiles_noqueue
LennyN95 Nov 21, 2024
b570e8a
dev-only: remove enforced num_processes = 1
LennyN95 Nov 21, 2024
3628055
add env to overwrite npp and nps
LennyN95 Nov 21, 2024
6bcec21
Fix preprocessor initialization order
surajpaib Nov 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
283 changes: 184 additions & 99 deletions nnunetv2/inference/data_iterators.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,44 @@ def preprocess_fromfiles_save_to_queue(list_of_lists: List[List[str]],
raise e


def preprocess_fromfiles_noqueue(list_of_lists: List[List[str]],
list_of_segs_from_prev_stage_files: Union[None, List[str]],
output_filenames_truncated: Union[None, List[str]],
plans_manager: PlansManager,
dataset_json: dict,
configuration_manager: ConfigurationManager,
verbose: bool = False):

print("Running preprocessing in non-multiprocessing mode")

data_iterator = []
label_manager = plans_manager.get_label_manager(dataset_json)
preprocessor = configuration_manager.preprocessor_class(verbose=verbose)

for idx in range(len(list_of_lists)):

input_files = list_of_lists[idx]
seg_file = list_of_segs_from_prev_stage_files[idx] if list_of_segs_from_prev_stage_files is not None else None
output_file = output_filenames_truncated[idx] if output_filenames_truncated is not None else None

data, seg, data_properties = preprocessor.run_case(input_files, seg_file, plans_manager, configuration_manager, dataset_json)

if list_of_segs_from_prev_stage_files is not None and list_of_segs_from_prev_stage_files[idx] is not None:
seg_onehot = convert_labelmap_to_one_hot(seg[0], label_manager.foreground_labels, data.dtype)
data = np.vstack((data, seg_onehot))

data = torch.from_numpy(data).to(dtype=torch.float32, memory_format=torch.contiguous_format)

preprocessed_data = {
'data': data,
'data_properties': data_properties,
'ofile': output_file
}

data_iterator.append(preprocessed_data)

return data_iterator

def preprocessing_iterator_fromfiles(list_of_lists: List[List[str]],
list_of_segs_from_prev_stage_files: Union[None, List[str]],
output_filenames_truncated: Union[None, List[str]],
Expand All @@ -67,56 +105,64 @@ def preprocessing_iterator_fromfiles(list_of_lists: List[List[str]],
num_processes: int,
pin_memory: bool = False,
verbose: bool = False):
context = multiprocessing.get_context('spawn')
manager = Manager()
num_processes = min(len(list_of_lists), num_processes)
assert num_processes >= 1
processes = []
done_events = []
target_queues = []
abort_event = manager.Event()
for i in range(num_processes):
event = manager.Event()
queue = Manager().Queue(maxsize=1)
pr = context.Process(target=preprocess_fromfiles_save_to_queue,
args=(
list_of_lists[i::num_processes],
list_of_segs_from_prev_stage_files[
i::num_processes] if list_of_segs_from_prev_stage_files is not None else None,
output_filenames_truncated[
i::num_processes] if output_filenames_truncated is not None else None,
plans_manager,
dataset_json,
configuration_manager,
queue,
event,
abort_event,
verbose
), daemon=True)
pr.start()
target_queues.append(queue)
done_events.append(event)
processes.append(pr)

worker_ctr = 0
while (not done_events[worker_ctr].is_set()) or (not target_queues[worker_ctr].empty()):
# import IPython;IPython.embed()
if not target_queues[worker_ctr].empty():
item = target_queues[worker_ctr].get()
worker_ctr = (worker_ctr + 1) % num_processes
else:
all_ok = all(
[i.is_alive() or j.is_set() for i, j in zip(processes, done_events)]) and not abort_event.is_set()
if not all_ok:
raise RuntimeError('Background workers died. Look for the error message further up! If there is '
'none then your RAM was full and the worker was killed by the OS. Use fewer '
'workers or get more RAM in that case!')
sleep(0.01)
continue
if pin_memory:
[i.pin_memory() for i in item.values() if isinstance(i, torch.Tensor)]
yield item
[p.join() for p in processes]

if num_processes > 1:
context = multiprocessing.get_context('spawn')
manager = Manager()
num_processes = min(len(list_of_lists), num_processes)
assert num_processes >= 1
processes = []
done_events = []
target_queues = []
abort_event = manager.Event()
for i in range(num_processes):
event = manager.Event()
queue = Manager().Queue(maxsize=1)
pr = context.Process(target=preprocess_fromfiles_save_to_queue,
args=(
list_of_lists[i::num_processes],
list_of_segs_from_prev_stage_files[
i::num_processes] if list_of_segs_from_prev_stage_files is not None else None,
output_filenames_truncated[
i::num_processes] if output_filenames_truncated is not None else None,
plans_manager,
dataset_json,
configuration_manager,
queue,
event,
abort_event,
verbose
), daemon=True)
pr.start()
target_queues.append(queue)
done_events.append(event)
processes.append(pr)

worker_ctr = 0
while (not done_events[worker_ctr].is_set()) or (not target_queues[worker_ctr].empty()):
if not target_queues[worker_ctr].empty():
item = target_queues[worker_ctr].get()
worker_ctr = (worker_ctr + 1) % num_processes
else:
all_ok = all(
[i.is_alive() or j.is_set() for i, j in zip(processes, done_events)]) and not abort_event.is_set()
if not all_ok:
raise RuntimeError('Background workers died. Look for the error message further up! If there is '
'none then your RAM was full and the worker was killed by the OS. Use fewer '
'workers or get more RAM in that case!')
sleep(0.01)
continue
if pin_memory:
[i.pin_memory() for i in item.values() if isinstance(i, torch.Tensor)]
yield item
[p.join() for p in processes]
else:
print("Running preprocessing in non-multiprocessing mode")
data_iterator = preprocess_fromfiles_noqueue(list_of_lists, list_of_segs_from_prev_stage_files, output_filenames_truncated, plans_manager, dataset_json, configuration_manager, verbose=verbose)
for item in data_iterator:
if pin_memory:
[i.pin_memory() for i in item.values() if isinstance(i, torch.Tensor)]
yield item


class PreprocessAdapter(DataLoader):
Expand Down Expand Up @@ -253,6 +299,37 @@ def preprocess_fromnpy_save_to_queue(list_of_images: List[np.ndarray],
raise e


def preprocess_fromnpy_noqueue(list_of_images: List[np.ndarray],
list_of_segs_from_prev_stage: Union[List[np.ndarray], None],
list_of_image_properties: List[dict],
truncated_ofnames: Union[List[str], None],
plans_manager: PlansManager,
dataset_json: dict,
configuration_manager: ConfigurationManager,
verbose: bool = False):
print("Running preprocessing in non-multiprocessing mode")
data_iterator = []
label_manager = plans_manager.get_label_manager(dataset_json)
preprocessor = configuration_manager.preprocessor_class(verbose=verbose)

for i in range(len(list_of_images)):
image = list_of_images[i]
seg_prev_stage = list_of_segs_from_prev_stage[i] if list_of_segs_from_prev_stage is not None else None
props = list_of_image_properties[i]
ofname = truncated_ofnames[i] if truncated_ofnames is not None else None
data, seg = preprocessor.run_case_npy(image, seg_prev_stage, props, plans_manager, configuration_manager, dataset_json)
if seg_prev_stage is not None:
seg_onehot = convert_labelmap_to_one_hot(seg[0], label_manager.foreground_labels, data.dtype)
data = np.vstack((data, seg_onehot))
preprocessed_data = {
'data': data,
'data_properties': props,
'ofile': ofname if ofname is not None else None
}
data_iterator.append(preprocessed_data)
return data_iterator


def preprocessing_iterator_fromnpy(list_of_images: List[np.ndarray],
list_of_segs_from_prev_stage: Union[List[np.ndarray], None],
list_of_image_properties: List[dict],
Expand All @@ -263,52 +340,60 @@ def preprocessing_iterator_fromnpy(list_of_images: List[np.ndarray],
num_processes: int,
pin_memory: bool = False,
verbose: bool = False):
context = multiprocessing.get_context('spawn')
manager = Manager()
num_processes = min(len(list_of_images), num_processes)
assert num_processes >= 1
target_queues = []
processes = []
done_events = []
abort_event = manager.Event()
for i in range(num_processes):
event = manager.Event()
queue = manager.Queue(maxsize=1)
pr = context.Process(target=preprocess_fromnpy_save_to_queue,
args=(
list_of_images[i::num_processes],
list_of_segs_from_prev_stage[
i::num_processes] if list_of_segs_from_prev_stage is not None else None,
list_of_image_properties[i::num_processes],
truncated_ofnames[i::num_processes] if truncated_ofnames is not None else None,
plans_manager,
dataset_json,
configuration_manager,
queue,
event,
abort_event,
verbose
), daemon=True)
pr.start()
done_events.append(event)
processes.append(pr)
target_queues.append(queue)

worker_ctr = 0
while (not done_events[worker_ctr].is_set()) or (not target_queues[worker_ctr].empty()):
if not target_queues[worker_ctr].empty():
item = target_queues[worker_ctr].get()
worker_ctr = (worker_ctr + 1) % num_processes
else:
all_ok = all(
[i.is_alive() or j.is_set() for i, j in zip(processes, done_events)]) and not abort_event.is_set()
if not all_ok:
raise RuntimeError('Background workers died. Look for the error message further up! If there is '
'none then your RAM was full and the worker was killed by the OS. Use fewer '
'workers or get more RAM in that case!')
sleep(0.01)
continue
if pin_memory:
[i.pin_memory() for i in item.values() if isinstance(i, torch.Tensor)]
yield item
[p.join() for p in processes]
if num_processes > 1:
context = multiprocessing.get_context('spawn')
manager = Manager()
num_processes = min(len(list_of_images), num_processes)
assert num_processes >= 1
target_queues = []
processes = []
done_events = []
abort_event = manager.Event()
for i in range(num_processes):
event = manager.Event()
queue = manager.Queue(maxsize=1)
pr = context.Process(target=preprocess_fromnpy_save_to_queue,
args=(
list_of_images[i::num_processes],
list_of_segs_from_prev_stage[
i::num_processes] if list_of_segs_from_prev_stage is not None else None,
list_of_image_properties[i::num_processes],
truncated_ofnames[i::num_processes] if truncated_ofnames is not None else None,
plans_manager,
dataset_json,
configuration_manager,
queue,
event,
abort_event,
verbose
), daemon=True)
pr.start()
done_events.append(event)
processes.append(pr)
target_queues.append(queue)

worker_ctr = 0
while (not done_events[worker_ctr].is_set()) or (not target_queues[worker_ctr].empty()):
if not target_queues[worker_ctr].empty():
item = target_queues[worker_ctr].get()
worker_ctr = (worker_ctr + 1) % num_processes
else:
all_ok = all(
[i.is_alive() or j.is_set() for i, j in zip(processes, done_events)]) and not abort_event.is_set()
if not all_ok:
raise RuntimeError('Background workers died. Look for the error message further up! If there is '
'none then your RAM was full and the worker was killed by the OS. Use fewer '
'workers or get more RAM in that case!')
sleep(0.01)
continue
if pin_memory:
[i.pin_memory() for i in item.values() if isinstance(i, torch.Tensor)]
yield item
[p.join() for p in processes]
else:
print("Running preprocessing in non-multiprocessing mode")
data_iterator = preprocess_fromnpy_noqueue(list_of_images, list_of_segs_from_prev_stage, list_of_image_properties, truncated_ofnames, plans_manager, dataset_json, configuration_manager, verbose=verbose)
for item in data_iterator:
if pin_memory:
[i.pin_memory() for i in item.values() if isinstance(i, torch.Tensor)]
yield item
Loading