有些服务器上可能安装了多块GPU供大家共同使用,为了不产生冲突,有时需要按情况指定自己的程序具体在哪些GPU上运行。
(注意:本文主要参考了https://blog.csdn.net/qq_38451119/article/details/81065675,我这里主要是记录一下,以后自己再用到的时候方便查找)
下面是实验室GPU的情况:
下面是具体的方法:
1. 在python代码中通过CUDA_VISIBLE_DEVICES来指定
比如,我要使用上面编号为“3”的GPU来运行我的程序,则需要在自己的程序中加入以下代码:
1 import os 2 os.environ('CUDA_VISIBLE_DEVICES') = '3'
下面可以通过命令:nvidia-smi 来查看GPU的使用情况:
可以看到编号为“3”的GPU已经被我占用了,而其他几块GPU则没被占用,这种就相当于在我们运行程序的时候,将除编号为“3”的GPU以外的其他GPU全部屏蔽了,只有编号为“3”的GPU对当前运行的程序是可见的。
同样,如果要占用多块GPU的话,比如使用编号为”1,2“的GPU来训练,则上面的代码可以改成:
import os os.environ('CUDA_VISIBLE_DEVICES')='1,2'
需要注意的是,仅做上述设置还不够,还需要在代码中进行一些设置。相较于pytorch,Keras等,TensorFlow的设置要麻烦一些。多GPU并行可分为模型并行和数据并行两大类,这里介绍一下数据并行方式。数据并行原理:模型参数保存在一个指定gpu/cpu上,模型参数的副本在不同gpu上,每次训练,提供batch_size*gpu_num数据,并等量拆分成多个batch,分别送入不同GPU。前向在不同gpu上进行,模型参数更新时,将多个GPU后向计算得到的梯度数据进行平均,并在指定GPU/CPU上利用梯度数据更新模型参数。下面以一个例子来说明:
1 #计算平均梯度的函数 2 def average_gradients(tower_grads): 3 average_grads=[] 4 for grad_and_vars in zip(*tower_grads): 5 grads=[] 6 for g,_ in grad_and_vars: 7 expend_g=tf.expand_dims(g,0) 8 grads.append(expend_g) 9 grad=tf.concat(grads,0) 10 grad=tf.reduce_mean(grad,0) 11 v=grad_and_vars[0][1] 12 grad_and_var=(grad,v) 13 average_grads.append(grad_and_var) 14 return average_grads 15 16 #设置GPU数量,每块GPU上运行的batch大小 17 num_gpus = 2 18 each_gpu_batch = 16 19 batch_size = each_gpu_batch * num_gpus 20 21 #占位符 22 data = tf.placeholder(tf.float32, [batch_size, img_size, img_size, 3]) 23 opt = tf.train.AdamOptimizer(learning_rate=d_lr, beta1=0.5, beta2=0.9) 24 25 #存储所有GPU上的梯度 26 tower_grads = [] 27 28 # 多GPU并行工作 29 with tf.variable_scope(tf.get_variable_scope()): 30 for i in range(num_gpus): 31 with tf.device("/gpu:%d" % i): 32 with tf.name_scope("tower_%d" % i): 33 _data = data[i*each_gpu_batch:(i+1)*each_gpu_batch] 34 generated = network(_data, s_size=s_size) 35 loss = loss(_data, generated) 36 37 grads = opt.compute_gradients(loss) 38 tower_grads.append(grads) 39 40 tf.get_variable_scope().reuse_variables() 41 42 grads = average_gradients(tower_grads) 43 train_opt = opt.apply_gradients(grads) 44 45 #后面就正常的调用session运行就好了
注意:上述代码并不能直接运行,只是给出了数据并行所需的步骤。详细解释参考:
https://blog.csdn.net/minstyrain/article/details/80986397
https://www.cnblogs.com/hrlnw/p/7779058.html
2. 在终端执行.py文件时通过CUDA_VISIBLE_DEVICES来指定
该方法与方法1本质上是一样的,只不过一个是在代码中指定GPU,一个是在执行代码的命令行中指定GPU。
命令行形式如下:
CUDA_VISIBLE_DEVICES=1 python train.py
3. 在python代码中通过tf.device()函数来指定训练时所要使用的GPU
比如,我要使用上面编号为“1”的GPU来运行我的程序,则需要在自己的程序中加入以下代码:
import tensorflow as tf tf.device('/gpu:1')
下面是实验结果(注:我自己没尝试这种方式,所以下面的实验结果图示直接用的博客https://blog.csdn.net/qq_38451119/article/details/81065675中的):
可以看到,这样指定GPU还是有一点毛病的。虽然指定了第“1“块GPU来训练,但是其它几个GPU也还是被占用,只是模型训练的时候,是在第1块GPU上进行。所以,最好使用前面介绍的第1,2种方法。