【技术分享】Python多进程处理，处理效率翻N倍

Python多进程处理，处理效率翻N倍

前言

多线程是个好东西，对于一大批需要处理的文件，如果只用一个进程跑的话，会花很多时间，但是，如果你将这些文件分批，分给不同的进程跑，处理速度将会大大加快。

流程

以MRI重建里面的数据处理为例，在MRI重建模型训练之前，你需要对全采样的图片进行欠采样，这个过程比较耗时，2000多张图片你单进程处理需要1-2个小时，这是十分耗时的，这个时候就可
以使用多线程，步骤如下：

初始化多线程

首先使用如下命令初始化多进程，如下：

import multiprocessing as mp
mp.set_start_method('spawn', force=True)

注意使用force=True，不然会报错：context has already been set

定义子进程处理函数

这个就是处理数据的关键部分，你把相关处理流程写在里面就行，模板如下：

import threading
lock = threading.Lock()
def sub_processor(pid, img_files, data_path, save_path):
    text = 'processor %d' % pid
    with lock:
        progress = tqdm(
            total=len(img_files),
            position=pid,
            desc=text,
            ncols=0
        )
    for img_name in img_files:
    	# 处理流程，这里省略
    	
    	# 更新进度
    	with lock:
            progress.update(1)
    with lock:
        progress.close()

同时这里使用了锁机制，更好地更新进度条，不需要可以去掉。

定义相关参数

首先定义你的线程数目，以及其他参数，如数据集地址，欠采样图片要保存地址：

# 全采样
data_path = "val_HQ_60percent_MRI"
# 欠采样
save_path = "val_LQ_60percent_MRI"
os.makedirs(save_path, exist_ok=True)
# 要处理的图片
img_files = sorted(os.listdir(data_path))
img_num = len(img_files)
# 进程数目
thread_num = 2

划分数据并分配给不同的进程

processes = []
img_num = len(img_files)
per_thread_img_num = img_num // thread_num
start_time = time.time()
print('Start process')
# 给每个进程分配处理的数据
for i in range(thread_num):
    if i == thread_num - 1:
        sub_img_list = img_files[i * per_thread_img_num:]
    else:
        sub_img_list = img_files[i * per_thread_img_num: (i + 1) * per_thread_img_num]
    p = mp.Process(target=sub_processor, args=(i, sub_img_list, data_path, save_path))
    p.start()
    processes.append(p)

for p in processes:
    p.join()

end_time = time.time()
total_time = end_time - start_time

print("Total process time: %.4f s" % (total_time))

注意mp.Process(target=sub_processor, args=(i, sub_img_list, data_path, save_path))中的args就是sub_processor函数要传递的参数，要对应。

最后，主函数运行要在main函数里面，即：

if __name__=='__main__':
	# 上面的处理代码

否则会报错：

An attempt has been made to start a new process before the  current process has finished its bootstrapping phase

完整代码

关注公众号世说CV回复“多进程”即可获得，没有回复可后台私信。