ICode9

精准搜索请尝试: 精确搜索
首页 > 其他分享> 文章详细

pytorch中使用cuda进行多任务multiprocessing

2021-02-16 23:30:41  阅读:4535  来源: 互联网

标签:spawn torch pytorch cuda context 进程 multiprocessing pool


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

感谢您的阅读!

标签:spawn,torch,pytorch,cuda,context,进程,multiprocessing,pool
来源: https://blog.csdn.net/YNNAD1997/article/details/113829532

本站声明: 1. iCode9 技术分享网(下文简称本站)提供的所有内容,仅供技术学习、探讨和分享;
2. 关于本站的所有留言、评论、转载及引用,纯属内容发起人的个人观点,与本站观点和立场无关;
3. 关于本站的所有言论和文字,纯属内容发起人的个人观点,与本站观点和立场无关;
4. 本站文章均是网友提供,不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属;如您发现该文章侵犯了您的权益,可联系我们第一时间进行删除;
5. 本站为非盈利性的个人网站,所有内容不会用来进行牟利,也不会利用任何形式的广告来间接获益,纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

专注分享技术,共同学习,共同进步。侵权联系[81616952@qq.com]

Copyright (C)ICode9.com, All Rights Reserved.

ICode9版权所有