如何在 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，所以我不想知道从操作系统内核获取 GPU 信息的方法。

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

原文由 Sangwon Kim 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

有一种未记录的方法称为 device_lib.list_local_devices() 使您能够列出本地进程中可用的设备。 （ 注意，作为一种未记录的方法，这会受到向后不兼容的更改的影响。）该函数返回 DeviceAttributes 协议缓冲区 对象的列表。您可以提取 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() 将运行一些初始化代码，默认情况下，将在所有设备上分配所有 GPU 内存（ GitHub 问题）。为避免这种情况，首先使用显式小的 per_process_gpu_fraction 或 allow_growth=True 创建一个会话，以防止分配所有内存。有关更多详细信息，请参阅 此问题。

原文由 mrry 发布，翻译遵循 CC BY-SA 3.0 许可协议

Answer 2

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

 from tensorflow.python.client import device_lib

device_lib.list_local_devices()

原文由 hyun woo Cho 发布，翻译遵循 CC BY-SA 3.0 许可协议