tensorflow2.x（二）多进程并行模型

tensorflow2.x（二）多进程并行模型

在上一篇文章中，我们解决了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算力。

运行nvidia-smi可看到如下，开启了6个线程，同时运行模型：