python for循环并行，单卡多卡并行解决办法

一、前言

当你显存足够，想用空间换时间；或者想单纯的加速模型。正好用的是python的for循环，而且每一次循环都是一次模型的推理过程，那么本 blog内容将非常适合你。

二、单卡并行

现在并行的方法有很多，本文没有对现有的方法进行归类，但是找到了适合自己用的，就分享在这里。比如:

from concurrent.futures.thread import ThreadPoolExecutor

def job(a, b):
	return a+b

def parallel_run(input_dir, a, b, num_worker):
    thread_pool = ThreadPoolExecutor(max_workers=num_worker)
    print('Thread Pool is created!')
    for i in range(len(input_dir)):
        thread_pool.submit(job, a, b)
    thread_pool.shutdown(wait=True)

再比如异步执行:

from multiprocessing import Pool

# job可以是你的推理函数
def job(a, b, c):
	return a+b+c

num_worker = 16  # 线程池数量
input_dir = ''
result_list = []  # 存储多线程的返回值
a, b = 1, 2
pool = Pool(num_worker)
for i, indis in enumerate(input_dir): 
    result_list.append(pool.apply_async(func=job, args=(a, i, b)))
pool.close()
pool.join()  # join的目的是把for循环的多线程都执行完，再执行后面的其他进程。

# 异步async的返回值需要通过get获取
res = []
for i in result_list:
    res.append(i.get())
# res即函数的返回值

上面的方法二选一，笔者使用的是第二种方法

三、多卡并行

直接修改job函数即可

def job(a, i, c):
    available_GPU = [0,1,2,3]  # 指定你可以使用的GPU_id
    GPU_ID = i % len(available_GPU)  # 将for循环中的每个推理按照顺序放在不同的GPU上
    torch.cuda.set_device(GPU_ID)  # 设置GPU
    a = torch.tensor(a).cuda()  # 将要处理的数据放在cuda上，原来是np格式，转成tensor可以直接用
    c = torch.tensor(c).cuda()  # 将要处理的数据放在cuda上，原来是np格式，转成tensor可以直接用
    return infer(a, c)  # infer为模型推理函数

四、踩的一些坑

1. # 记得不要把.get()接到异步函数后面，像这样会拖慢并行的速度，各位bloger可以自行尝试
    result_list.append(pool.apply_async(func=job, args=(a, i, b)).get())
2. # 在异步调用class中的函数时，千万不要用装饰器，否则会提示
    Can't pickle __main__.Test'>: it's not the same object as __main__.Test

五、小结

本文仅简单记录笔者学到的方法和踩到的坑，欢迎各位指点！~