pytorch中使用cuda进行多任务multiprocessing

我们都知道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 with ctx.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

感谢您的阅读！