在做一个图像处理相关的项目,NVIDIA基于CUDA的GPU,我们需要把图像处理做成服务,并提升CPU上并发速度。整个服务既有CPU处理,又有GPU处理,我们就需要把CPU上的处理做成多线并发,把GPU上的数据做成batch并发起来。由于code是用pytorch 的python版本实现的,而不是c++,这就给我们造成了困扰,对于python我们知道多进程才能做到利用CPU多核的目的,而多线并不能,我们这里采用多进程进行并发。而用多进程进行并发的时候,就出现了如下的error。
CUDA error (3): initialization error
经过查找资料,pyTorch的讨论区给出了解释,就是我们在执行多进程之前进行了CUDA相关的操作,造成启动的进程并没有获得CUDA初始化的结果。而且我们测试发现,即使我们在global区域进行初始化CUDA的操作,在每个进程依然会调用。我们要做的就是把CUDA的初始化部分放到每个进程里面,然后在main函数的第一句加上
mp.set_start_method('spawn')
修改后的主函数的代码片段如下:
....
....
....
if __name__ == "__main__":
# to make sure multiprocess running on CUDA, you have to set start method as "spawn"
mp.set_start_method('spawn')
manager = mp.Manager()
q = manager.Queue()
lock = manager.Lock()
gpuLock = manager.Lock()
p = Pool()
p.apply_async(serverSocket, args=(q, lock))
# start 10 processor
for i in range(3):
p.apply_async(run_task, args=(q, lock, gpuLock))
p.close()
p.join()