pytorch 指定卡1_如何为TensorFlow和PyTorch自动选择空闲GPU，解决抢卡争端

原标题：如何为TensorFlow和PyTorch自动选择空闲GPU，解决抢卡争端

雷锋网按：本文作者天清，原文载于其知乎专栏世界那么大我想写代码，雷锋网获其授权发布。

项目地址：QuantumLiu/tf_gpu_manager

***

更新：支持pytorch

***

使用

git clone https://github.com/QuantumLiu/tf_gpu_manager

把manager.py放到你训练的目录就行。

直接使用with gm.auto_choice()自动选择设备进行接下来代码块的操作。

import tensorflow as tf

from manager import GPUManager

from keras.layers LSTM

gm=GPUManager()

with gm.auto_choice():

x=tf.placeholder(tf.float32,shape=(None,20,64))

y=LSTM(32)(x)背景

随着深度学习技术快速的发展，深度学习任务的数据和计算规模也越来越大，想要做出个像样的work，没有一台powerful的GPU工作站是万万不能的。

除了要求单卡性能强大，GPU数量多也很重要。

因为以下几点原因，多GPU工作站已经成了各大实验室的标配：

一般来说，一个深度学习项目需要一个实验室或者小组的多人合作完成，要共享一台或几台工作站。一个host多个GPU比较方便。

实验需要试多组参数或者对比试验。多GPU并行跑省时间。

模型计算量大，需要将模型不同分配在多个GPU上计算。

现在，Tensorflow、pytorch等主流深度学习框架都支持多GPU训练。

比如Tensorflow，在 tensorflow\python\framework中定义了device函数，返回一个用来执行操作的GPU设备的context manager对象。

def device(device_name_or_function):

"""Wrapper for `Graph.device()` using the default graph.

See

@{tf.Graph.device}

for more details.

Args:

device_name_or_function: The device name or function to use in the context.

Returns:

A context manager that specifies the default device to use for newly created ops.

"""

return get_default_graph().device(device_name_or_function)

在我们的训练脚本中使用with语句就可以指定接下来的操作在某个GPU上进行。

with tf.device('/gpu:2'):

a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')

b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')

c = tf.matmul(a, b)

那么问题来了：

在写训练脚本时怎么知道哪个GPU是空闲可用的？

同组的人做实验和我冲突怎么办？

将来某个时刻运行这个脚本的时候是不是还要根据情况修改？

同行用我的代码复现实验，GPU配置环境不一样，他们甚至可能没有GPU，又要改代码？

当然，上道儿的开发者都知道nvidia-smi可以查询显卡信息，查看GPU显存、温度、功率使用，然后选择合适的GPU。

<img src="https://static.leiphone.com/uploads/new/article/pic/201708/925d4b6331f0d6a55b25b9598bf86fc0.png" data-rawwidth="1109" data-rawheight="727" class="origin_image zh-lightbox-thumb" width="1109" data-original="https://pic4.zhimg.com/v2-4e889c4ab7942cc4056249f2f063f28f_r.png" _src="https://static.leiphone.com/uploads/new/article/pic/20