python 并行化处理多模型多GPU

python 并行化处理多模型多GPU

在基于深度网络的各种应用中，串行地遍历超参数组合完成参数优化速度缓慢。在人为指定多个参数组合情况下，使用多核多卡，开并行能够显著加快模型的运行速度。以下为本人使用的一些代码。

import os
import torch
import multiprocessing as mp
from multiprocessing.pool import Pool  # 进程池
from multiprocessing import cpu_count  #获取系统默认CPU内核数目

def one_model(para):
    para1 = para[0]
    para2 = para[1]
    GPU = para[2]

    os.environ["CUDA_VISIBLE_DEVICES"] = str(GPU)
    torch.cuda.set_device(GPU)
    model.run(para1,para2)
    return model.res

def main():
	mp.set_start_method('spawn') #对于多GPU分别运行不同的模型时是必要的
    results = []
    p = Pool(processes=cpu_count())
    res = p.map(one_model, ([(para1,para2,GPU_num) for num,i in zip(para1_list, para2_list,GPU_list)]))
    p.close()
    p.join()

    for item in res:
        results.append(item)
main()