tensorflow显存管理、tensorflow使用多个gpu训练

 通常在程序开始之前并不知道需要多大的显存,程序会去申请GPU的显存的50%

比如一个8G的内存,被占用了2G,那么程序会申请4G的显存(因为有足够的剩余显存)

如果此时显存被占用7G,那么程序会申请剩下的所有的1G的显存。也许你的程序根本用不着这么多显存,200M也许就够了,这时候如果程序能按照需求去申请就好了,幸运的是,这样的方法是存在的:

import tensorflow as tf  
import os  
os.environ["CUDA_VISIBLE_DEVICES"] = '0'   #指定第一块GPU可用  
config = tf.ConfigProto()  
config.gpu_options.per_process_gpu_memory_fraction = 0.5  # 程序最多只能占用指定gpu50%的显存  
config.gpu_options.allow_growth = True      #程序按需申请内存  
sess = tf.Session(config = config)  

 

能看到,只使用单个GPU跑程序,但三块显卡的显存都被占用。

这是因为TensorFlow训练时默认占用所有GPU的显存。

这样如果有人想使用其他两个GPU跑程序,就会因为显存不足而无法运行。 
所以需要人为指定显存占用率或指定使用单张显卡。

根据 TF官网tutorial部分的Using GPUs,可以总结三种方法:

  1. 第一种是使用 allow_growth,实现显存运行时分配。当allow_growth设置为True时,TF程序一开始被分配很少的显存,在运行时根据需求增长而扩大显存的占用。

    config = tf.ConfigProto()  
    config.gpu_options.allow_growth = True  
    sess = tf.Session(config=config, ...)  
  2. 第二种是使用 per_process_gpu_memory_fraction,指定每个可用GPU的显存分配率。在构造tf.Session()时候通过tf.GPUOptions配置参数,显式地指定需要分配的显存比例。

    
    #告诉TF它可以使用每块GPU显存的40%  
    
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = 0.4
    session = tf.Session(config=config, ...)

    这种方法指定了每个GPU进程的显存占用上限,但它会同时作用于所有GPU,不能对不同GPU设置不同的上限。

  3. 在运行训练程序前,在用户根目录下配置环境(~/.bashrc):

    export CUDA_VISIBLE_DEVICES = NUM  

    NUM是用户指定显卡的序号(0,1,2…),可以先用 nvidia-smi 查看当前哪块显卡可用。但这种方法限制了用户可见的GPU数量,比如你的其他程序在你的目录里无法选择别的GPU; 你的程序也没法使用multiple GPUs。

tensorflow使用多个gpu训练

https://www.cnblogs.com/hrlnw/p/7779058.html

关于多gpu训练,tf并没有给太多的学习资料,比较官方的只有

:tensorflow-models/tutorials/image/cifar10/cifar10_multi_gpu_train.py

但代码比较简单,只是针对cifar做了数据并行的多gpu训练,利用到的layer、activation类型不多,针对更复杂网络的情况,并没有给出指导。

一、思路

单GPU时,思路很简单,前向、后向都在一个GPU上进行,模型参数更新时只涉及一个GPU。多GPU时,有模型并行和数据并行两种情况。模型并行指模型的不同部分在不同GPU上运行。数据并行指不同GPU上训练数据不同,但模型是同一个(相当于是同一个模型的副本)。在此只考虑数据并行,这个在tf的实现思路如下:

模型参数保存在一个指定gpu/cpu上,模型参数的副本在不同gpu上,每次训练,提供batch_size*gpu_num数据,并等量拆分成多个batch,分别送入不同GPU。前向在不同gpu上进行,模型参数更新时,将多个GPU后向计算得到的梯度数据进行平均,并在指定GPU/CPU上利用梯度数据更新模型参数。

假设有两个GPU(gpu0,gpu1),模型参数实际存放在cpu0上,实际一次训练过程如下图所示:

二、tf代码实现

大部分需要修改的部分集中在构建计算图上,假设在构建计算图时,数据部分基于tensorflow1.4版本的dataset类,那么代码要按照如下方式编写:

next_img, next_label = iterator.get_next()
image_splits = tf.split(next_img, num_gpus)
label_splits = tf.split(next_label, num_gpus)
tower_grads = []
tower_loss = []
counter = 0
for d in self.gpu_id:
    with tf.device('/gpu:%s' % d):
        with tf.name_scope('%s_%s' % ('tower', d)):
            cross_entropy = build_train_model(image_splits[counter], label_splits[counter], for_training=True)
            counter += 1
            with tf.variable_scope("loss"):
                grads = opt.compute_gradients(cross_entropy)
                tower_grads.append(grads)
                tower_loss.append(cross_entropy)
                tf.get_variable_scope().reuse_variables()

mean_loss = tf.stack(axis=0, values=tower_loss)
mean_loss = tf.reduce_mean(mean_loss, 0)
mean_grads = util.average_gradients(tower_grads)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
    train_op = opt.apply_gradients(mean_grads, global_step=global_step)

第1行得到image和对应label

第2-3行对image和label根据使用的gpu数量做平均拆分(默认两个gpu运算能力相同,如果gpu运算能力不同,可以自己设定拆分策略)

第 4-5行,保存来自不同GPU计算出的梯度、loss列表

第7-16行,开始在每个GPU上创建计算图,最重要的是14-16三行,14,15把当前GPU计算出的梯度、loss值append到列表后,以便后续计算平均值。16行表示同名变量将会复用,这个是什么意思呢?假设现在gpu0上创建了两个变量var0,var1,那么在gpu1上创建计算图的时候,如果还有var0和var1,则默认复用之前gpu0上的创建的那两个值。

第18-20行计算不同GPU获取的grad、loss的平均值,其中第20行使用了cifar10_multi_gpu_train.py中的函数。

第23行利用梯度平均值更新参数。

 

注意:上述代码中,所有变量(vars)都放在了第一个GPU上,运行时会发现第一个GPU占用的显存比其他GPU多一些。如果想把变量放在CPU上,则需要在创建计算图时,针对每层使用到的变量进行设备指定,很麻烦,所以建议把变量放在GPU上。

单机多GPU训练

先简单介绍下单机的多GPU训练,然后再介绍分布式的多机多GPU训练。 
单机的多GPU训练, tensorflow的官方已经给了一个cifar的例子,已经有比较详细的代码和文档介绍, 这里大致说下多GPU的过程,以便方便引入到多机多GPU的介绍。 
单机多GPU的训练过程:

  1. 假设你的机器上有3个GPU;

  2. 在单机单GPU的训练中,数据是一个batch一个batch的训练。 在单机多GPU中,数据一次处理3个batch(假设是3个GPU训练), 每个GPU处理一个batch的数据计算。

  3. 变量,或者说参数,保存在CPU上

  4. 刚开始的时候数据由CPU分发给3个GPU, 在GPU上完成了计算,得到每个batch要更新的梯度。

  5. 然后在CPU上收集完了3个GPU上的要更新的梯度, 计算一下平均梯度,然后更新参数。

  6. 然后继续循环这个过程。

通过这个过程,处理的速度取决于最慢的那个GPU的速度。如果3个GPU的处理速度差不多的话, 处理速度就相当于单机单GPU的速度的3倍减去数据在CPU和GPU之间传输的开销,实际的效率提升看CPU和GPU之间数据的速度和处理数据的大小。

通俗解释

写到这里觉得自己写的还是不同通俗易懂, 下面就打一个更加通俗的比方来解释一下:

老师给小明和小华布置了10000张纸的乘法题并且把所有的乘法的结果加起来, 每张纸上有128道乘法题。 这里一张纸就是一个batch, batch_size就是128. 小明算加法比较快, 小华算乘法比较快,于是小华就负责计算乘法, 小明负责把小华的乘法结果加起来 。 这样小明就是CPU,小华就是GPU.

这样计算的话, 预计小明和小华两个人得要花费一个星期的时间才能完成老师布置的题目。 于是小明就招来2个算乘法也很快的小红和小亮。 于是每次小明就给小华,小红,小亮各分发一张纸,让他们算乘法, 他们三个人算完了之后, 把结果告诉小明, 小明把他们的结果加起来,然后再给他们没人分发一张算乘法的纸,依次循环,知道所有的算完。

这里小明采用的是同步模式,就是每次要等他们三个都算完了之后, 再统一算加法,算完了加法之后,再给他们三个分发纸张。这样速度就取决于他们三个中算乘法算的最慢的那个人, 和分发纸张的速度。

 

Guess you like

Origin blog.csdn.net/liulina603/article/details/80180355