如何在张量流中获取当前可用的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信息的方法。

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

Answer 1

您可以使用以下代码检查所有设备列表：from tensorflow.python.client import device_libdevice_lib.list_local_devices()

Answer 2

除了Mrry的出色解释之外，他建议在哪里使用，device_lib.list_local_devices()我可以向您展示如何从命令行检查GPU相关信息。因为目前只有Nvidia的GPU适用于NN框架，所以答案只涵盖了它们。Nvidia上有一个页面，其中记录了如何使用/ proc文件系统接口来获取有关驱动程序，任何已安装的NVIDIA图形卡以及AGP状态的运行时信息。/proc/driver/nvidia/gpus/0..N/information提供有关每个已安装的NVIDIA图形适配器的信息（型号名称，IRQ，BIOS版本，总线类型）。请注意，BIOS版本仅在X运行时可用。因此，您可以从命令行运行此命令，cat /proc/driver/nvidia/gpus/0/information并查看有关第一个GPU的信息。从python运行它很容易，并且您可以检查第二，第三，第四GPU直到失败。肯定地Mrry的答案更可靠，而且我不确定我的答案是否可以在非Linux机器上使用，但是Nvidia的页面提供了其他有趣的信息，但鲜为人知。