python多进程实现图片处理,进程池的大坑!
最近做大批量图片数据处理的时候,cpu占用一直很低,所以想用python多进程实现大量图片处理的加速。
但参考了很多很多博客,就是会有各种bug:跳过函数,函数不执行;只有主进程在运行,子进程都sleep等等。。。一一解决之后仍然跑不通,一行一行debug,找了n个小时终于发现了这个坑。
我们都知道python多进程是不共享数据的!
所以如果是用单写config的文件的,要记得在多进程传参的时候,额外把config文件的数据也传进去!
其实是知道这个点的,但之前就是没注意到这个问题,而且你只要之前定义过了变量在多进程的work_func()中也不会报错还能连接到那个数据,所以这个简单的坑困扰了很久。。
顺带一提,如果有device等等配置也需要传入哦~~
上代码:
import yaml
import os
import multiprocessing
# 图像处理
def process_image(img_path, config):
'''
图像处理操作
'''
if __name__=='__main__':
with open('./config.yaml', 'r', encoding='utf8') as fin:
config = yaml.safe_load(fin)
if not os.path.exists(config['data']['dst_label_path']):
os.makedirs(config['data']['dst_label_path'])
if not os.path.exists(config['data']['dst_image_path']):
os.makedirs(config['data']['dst_image_path'])
image_list = os.listdir(config['data']['src_image_path'])
image_list = [os.path.join(config['data']['src_image_path'], fname) for fname in image_list]
device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = MagicPoint(config['model'], input_channel=1, grid_size=8, device=device)
net.to(device).eval()
lock = multiprocessing.Lock()
# 创建进程池
pool = multiprocessing.Pool(processes=4)
for idx, fpath in tqdm(enumerate(image_list)):
# 多进程执行
pool.apply(process_image, args=(img_path, config))
# pool.apply_async(process_image(idx, fpath))
pool.close()
pool.join()
print('Done')