在CUDA上使用多进程

在CUDA上使用多进程

在批量处理大量数据时，我们常常会选择使用多进程来提高代码运行速度，充分利用CPU/GPU。

在此，记录了如何在CUDA上利用多进程来批量推理。

import multiprocessing as mp

# main下执行
torch.multiprocessing.set_start_method('forkserver',force = True)#此行代码必不可少

### CUDA下利用DL model进行推理
# 创建n个进程运行推理代码
pool_num = 4
divide_dir2dirs(pool_num, call_args['inputs'])  # 将一个目录下的所有图片均分到n个目录下
ctx = torch.multiprocessing.get_context("spawn")
pool = ctx.Pool(pool_num)
for ip in range(pool_num):
    tmp = copy.deepcopy(call_args)
    tmp['inputs'] += ('/dir'+str(ip+1))
    print(tmp)
    pool.apply_async(func=inference_worker, args=