系统相关
首页 > 系统相关> > tensorflow2.x(二)充分利用GPU资源、多进程并行模型

tensorflow2.x(二)充分利用GPU资源、多进程并行模型

作者:互联网

在上一篇文章中,我们解决了tensorflow在大样本训练中内存不足的问题,但是可能无法最大化利用GPU资源,在这篇文章中,我们继续挖掘如何充分利用GPU资源,将显卡的内存、算力全部拉满。

为了进一步挖掘显卡性能,进一步提升资源利用率,进一步解放双手,在这篇文章中,我们试图使用多进程,分配不同的显卡资源给多个模型,同时并行训练多个模型。
(虽然tf官方也有将显卡并行,使用多张显卡资源来提升计算效率的解决方案,但是仍然需要多源模型代码进行相关改变,且在笔者当下的尝试中,这种并行在生成环境中并不稳定)
因此,这里我们提出,同时并行多个模型的解决方案。

废话到此为止,直接上代码:

gpus = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
logical_gpus = tf.config.experimental.list_logical_devices('GPU')

这段代码是我前文中给出过的,也是tf官网里或是各种教程里轻而易举可以找到的。意思是,先获取你电脑中的全部可用GPU:gpus,然后给你当前进程中分配可见的gpus[0],再然后给当前可见的的gpus[0]分配可用内存memory_limit=1024。经过以上代码,你的当前进程中,就只能看到一张显卡,且显卡内存为memory_limit=1024

之后,通过多进程操作,我们就可以同时虚拟出多张显卡,每张显卡内独立跑各自的模型。

# 初始化cuda环境
def initial_cuda(gpu_name, memory_limit):
    gpus = tf.config.experimental.list_physical_devices('GPU')
    tf.config.experimental.set_visible_devices(gpus[gpu_name], 'GPU')
    tf.config.experimental.set_virtual_device_configuration(
        gpus[gpu_name],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=memory_limit)])
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')

# 配置当前进程中的cuda
def get_cuda(use_cuda):
    gpu_name, memory_limit, name = use_cuda
    initial_cuda(gpu_name, memory_limit)
    print("USE CUDA:", gpu_name, memory_limit, name)

# 根据显卡内存容量,每次开6个线程,结束后全部关闭,重新开,防止cuda初始化失败
def allocate_mp(param_all, func, cuda_list):
    mp_num = 6
    for i in range((int(len(param_all) / mp_num) + 1)):
        param_group = param_all[i * mp_num:(i + 1) * mp_num]
        for v in range(len(param_group)):
            param_group[v] = param_group[v] + (cuda_list[v], )
        pool = mp.Pool(processes=mp_num)
        pool.map(func, param_group)
        pool.close()

# 需要执行的tf函数
def func(parma):
	a,b,use_cuda = parma
    get_cuda(use_cuda)
    
# 根据自己情况,给GPU分组
cuda_list = list(itertools.product([0, 1], [2048], ['a', 'b', 'c']))
# 参数组
param_all = [(1,2),(3,4),(5,6)]
# 执行多线程
allocate_mp(param_all, func, cuda_list)

**对于cuda_list的解释:**这里我有[0,1]两张GPU,每个GPU分为[‘a’, ‘b’, ‘c’]三块,每块分配2048M内存。这样我可以得到6个虚拟GPU,支持6进程。

这里需要注意的一个坑在多进程中,进程池一旦打开,比如开2个进程,跑10组不同参数的函数。所有的函数是都在这固定的2个进程中进行,并不会在一个函数运行结束后关闭当前进程,然后重新开一个进程跑下一个函数。
因此,在我们的allocate_mp函数中,写了一个循环,对于超出6个参数组的情况,我们需要等上6个进程全部结束,然后再重新开新的6个进程,运行下一组。因为,在进程开启后,tensorflow的cuda初始化仅第一次有效,第二次执行带有cuda初始化的函数便会报错。

至此,就可以愉快地并行运行tensorflow模型,充分挖掘GPU算力。

标签:tensorflow2,显卡,list,并行,cuda,gpus,tf,GPU
来源: https://blog.csdn.net/qq_37373209/article/details/122756973