python多显卡支持_在多GPU情况下TensorFlow如何指定在哪些GPU上运行程序

最新推荐文章于 2023-11-16 10:11:27 发布

weixin_39887846

最新推荐文章于 2023-11-16 10:11:27 发布

阅读量735

点赞数

文章标签： python多显卡支持

有些服务器上可能安装了多块GPU供大家共同使用，为了不产生冲突，有时需要按情况指定自己的程序具体在哪些GPU上运行。

下面是实验室GPU的情况：

下面是具体的方法：

1. 在python代码中通过CUDA_VISIBLE_DEVICES来指定

比如，我要使用上面编号为“3”的GPU来运行我的程序，则需要在自己的程序中加入以下代码：

1 importos2 os.environ('CUDA_VISIBLE_DEVICES') = '3'

下面可以通过命令：nvidia-smi 来查看GPU的使用情况：

可以看到编号为“3”的GPU已经被我占用了，而其他几块GPU则没被占用，这种就相当于在我们运行程序的时候，将除编号为“3”的GPU以外的其他GPU全部屏蔽了，只有编号为“3”的GPU对当前运行的程序是可见的。

同样，如果要占用多块GPU的话，比如使用编号为”1，2“的GPU来训练，则上面的代码可以改成：

importos

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 defaverage_gradients(tower_grads):3 average_grads=[]4 for grad_and_vars in zip(*tower_grads):5 grads=[]6 for g,_ ingrad_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 returnaverage_grads15

16 #设置GPU数量，每块GPU上运行的batch大小

17 num_gpus = 2

18 each_gpu_batch = 16

19 batch_size = each_gpu_batch *num_gpus20

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 inrange(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运行就好了

注意：上述代码并不能直接运行，只是给出了数据并行所需的步骤。详细解释参考：

2. 在终端执行.py文件时通过CUDA_VISIBLE_DEVICES来指定

该方法与方法1本质上是一样的，只不过一个是在代码中指定GPU，一个是在执行代码的命令行中指定GPU。

命令行形式如下：

CUDA_VISIBLE_DEVICES=1 python train.py

3. 在python代码中通过tf.device()函数来指定训练时所要使用的GPU

比如，我要使用上面编号为“1”的GPU来运行我的程序，则需要在自己的程序中加入以下代码：

importtensorflow as tf