如何在tensorflow中获得当前可用的GPU？

Question

我有一个使用分布式TensorFlow的计划，我看到TensorFlow可以使用GPU进行训练和测试。在集群环境中，每台机器可以有0个或1个或多个GPU，我想在尽可能多的机器上将我的TensorFlow图运行到GPU中。

我发现在运行tf时。会话（）TensorFlow在日志消息中提供有关GPU的信息，如下所示：

I tensorflow/core/common_runtime/gpu/gpu_init.cc:126] DMA: 0 
I tensorflow/core/common_runtime/gpu/gpu_init.cc:136] 0:   Y 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:838] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GTX 1080, pci bus id: 0000:01:00.0)

我的问题是如何从TensorFlow获取当前可用GPU的信息？我可以从日志中获取加载的GPU信息，但我想用一种更复杂、更编程的方式。我还可以使用CUDA\u VISIBLE\u DEVICES环境变量有意限制GPU，因此我不想知道从操作系统内核获取GPU信息的方法。

简而言之，我想要一个类似于tf的函数。get\u available\u gpus（），如果机器中有两个可用的gpu，则返回['/gpu:0'，'/gpu:1']。我如何实现这一点？

Answer 1

测试工具中还有一个方法。因此，需要做的是：

tf.test.is_gpu_available()

和/或

tf.test.gpu_device_name()

在Tensorflow文档中查找参数。

Answer 2

您可以使用以下代码检查所有设备列表：

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

Answer 3

有一种未记录的方法称为device\u lib。list\u local\u devices（），用于列出本地进程中可用的设备。（注意，作为一种未记录的方法，这可能会发生向后不兼容的更改。）该函数返回设备属性列表。您可以提取GPU设备的字符串设备名称列表，如下所示：

from tensorflow.python.client import device_lib

def get_available_gpus():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos if x.device_type == 'GPU']

请注意（至少达到TensorFlow 1.4），调用设备库。list\u local\u devices（）将运行一些初始化代码，默认情况下，这些代码将在所有设备上分配所有GPU内存（GitHub问题）。为了避免这种情况，首先创建一个会话，其中每个进程的gpu分数明显较小，或者允许增长=True，以防止分配所有内存。有关更多详细信息，请参阅此问题。