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

Question

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

我发现在运行tf.Session()时，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_VISIBLE_DEVICES环境变量来限制GPU，所以我不想知道从OS内核获取GPU信息的方法。

简而言之，我想要一个类似tf.get_available_gpus()的函数，如果机器中有两个可用的gpu，它将返回['/gpu:0', '/gpu:1']。我该如何实现？

Answer 1

在test util中还有一个方法。 所以我们要做的就是：

tf.test.is_gpu_available()

和/或

tf.test.gpu_device_name()

查找Tensorflow文档中的参数。

Answer 2

有一个名为^{}的未记录方法，它使您能够列出本地进程中可用的设备。（N.B.作为一个未记录的方法，这会受到向后不兼容的更改的影响。）函数返回一个^{} protocol buffer对象列表。您可以提取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之前），调用device_lib.list_local_devices()将运行一些初始化代码，默认情况下，这些代码将分配所有设备（GitHub issue）上的所有GPU内存。为了避免这种情况，首先创建一个具有显式小的per_process_gpu_fraction或allow_growth=True的会话，以防止分配所有内存。有关详细信息，请参见this question。

Answer 3

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

from tensorflow.python.client import device_lib

device_lib.list_local_devices()