Coverage for src/image_utils/process_utils.py: 56%
34 statements
« prev ^ index » next coverage.py v7.2.5, created at 2023-05-19 16:15 -0700
« prev ^ index » next coverage.py v7.2.5, created at 2023-05-19 16:15 -0700
1from functools import partial
2from multiprocessing import Pool
3import random
4import time
5from pathlib import Path
6from joblib import dump, load
7import time
8import numpy as np
9from joblib import Parallel, delayed
10from joblib import Memory
13def timer_func(func):
14 def wrap_func(*args, **kwargs):
15 t1 = time.time()
16 result = func(*args, **kwargs)
17 t2 = time.time()
18 print(f'Function {func.__name__!r} executed in {(t2-t1):.4f}s')
19 return result
20 return wrap_func
23def split_range(a, n):
24 k, m = divmod(len(a), n)
25 return (a[i * k + min(i, m): (i + 1) * k + min(i + 1, m)] for i in range(n))
28def parallelize(func=None, num_processes=1, shuffle=False, custom_chunking=False, use_joblib=False):
29 def wrapper(iterable, **kwargs):
30 if use_joblib:
31 executor = Parallel(n_jobs=num_processes, max_nbytes=1e6)
32 tasks = (delayed(func)(iterable, idx, **kwargs) for idx in range(len(iterable)))
33 return executor(tasks)
34 else:
35 with Pool(num_processes) as p:
36 num_arr = list(range(len(iterable)))
37 if shuffle:
38 random.shuffle(num_arr)
39 for _ in p.imap_unordered(
40 partial(func, **kwargs),
41 split_range(num_arr, num_processes) if custom_chunking else iterable
42 ):
43 pass
45 return wrapper