我们都知道python有自带的multiprocessing模块,但是如果要使用cuda的话会报错:
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
但是查找torch中spawn,查找到torch.multiprocessing.
spawn,所得的介绍并不多而且网上搜到的都是抄这个说明书的
仔细阅读torch.multiprocessing的英文解释发现这个部分就是把python的官方multiprocessing给wrap(包)了一下,所以之前的应该都能用,因此我之前的pool代码可以直接使用
原来spawn的方法只是一种多任务的方法
spawn
父进程启动一个新的Python解释器进程。子进程只会继承那些运行进程对象的 run()
方法所需的资源。特别是父进程中非必须的文件描述符和句柄不会被继承。相对于使用 fork 或者 forkserver,使用这个方法启动进程相当慢。
只需要设定context为"spawn”即可,但是,如果直接使用
torch.multiprocessing.set_start_method("spawn")
会出bug:
RuntimeError('context has already been set')
查找stackflow中可以找到一位仁兄,跟我做的事情还差不多,都是pytorch想要多任务
https://github.com/pytorch/pytorch/issues/3492
从里边我使用的一种方法成功了,就是这个
I found a solution, which is to use a context object in multiprocessing.
ctx = multiprocessing.get_context("spawn")
And then replace
multiprocessing.xxx
withctx.xxx
.
我最终的代码为
def func(a,b,c):
return {"a":a,"b":b,"c":c}
def func2():
ctx = torch.multiprocessing.get_context("spawn")
print(torch.multiprocessing.cpu_count())
pool = ctx.Pool(5) # 7.7G
pool_list = []
for XX:
for XXX:
for XXX:
res = pool.apply_async(func, args=(a,b,c))
pool_list.append(res)
pool.close() # 关闭进程池,不再接受新的进程
pool.join() # 主进程阻塞等待子进程的退出
ccc = pd.DataFrame(columns=["a", "b", "c"])
for i in pool_list:
data = i.get()
ccc = ccc.append(data, ignore_index=True)
只是一个示意,供大家参考!
注意应该合理选择pool的worker个数,我这里是五个,然后使用了大约7.7G的显存。如果直接使用cpu_count()就是24个,显存直接爆掉。在实验的时候要使用如下代码来监视gpu使用情况:
watch -n 0.1 nvidia-smi
感谢您的阅读!