如何在张量流中获取当前可用的GPU？

我有一个使用分布式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']。我该如何实施？

有一个未记录的方法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 issue）。为避免这种情况，请首先使用一个显着小的per_process_gpu_fraction或创建一个会话allow_growth=True，以防止分配所有内存。有关更多详细信息，请参见此问题。

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

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

测试工具中还有一种方法。因此，所有要做的就是：

tf.test.is_gpu_available()

和/或

tf.test.gpu_device_name()

在Tensorflow文档中查找参数。

在TensorFlow 2.0中，您可以使用 tf.config.experimental.list_physical_devices('GPU')：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    print("Name:", gpu.name, "  Type:", gpu.device_type)

如果您安装了两个GPU，它将输出以下内容：

Name: /physical_device:GPU:0   Type: GPU
Name: /physical_device:GPU:1   Type: GPU

从2.1开始，您可以 experimental：

gpus = tf.config.list_physical_devices('GPU')

看到：

• 指南页
• 当前的API

该接受的答案给你GPU的数量，但它也分配所有这些GPU的内存。您可以通过在调用device_lib.list_local_devices（）之前创建具有固定较低内存的会话来避免这种情况，这对于某些应用程序可能是不需要的。

我最终使用nvidia-smi来获取GPU的数量，而没有在其上分配任何内存。

import subprocess

n = str(subprocess.check_output(["nvidia-smi", "-L"])).count('UUID')

除了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的页面提供了其他有趣的信息，但鲜为人知。

以下工作在tensorflow 2中：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    print("Name:", gpu.name, "  Type:", gpu.device_type)

从2.1开始，您可以删除experimental：

    gpus = tf.config.list_physical_devices('GPU')

https://www.tensorflow.org/api_docs/python/tf/config/list_physical_devices

NVIDIA GTX GeForce 1650 Ti我的机器中调用了一个GPUtensorflow-gpu==2.2.0

运行以下两行代码：

import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

输出：

Num GPUs Available:  1

使用这种方式并检查所有零件：

from __future__ import absolute_import, division, print_function, unicode_literals

import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_datasets as tfds


version = tf.__version__
executing_eagerly = tf.executing_eagerly()
hub_version = hub.__version__
available = tf.config.experimental.list_physical_devices("GPU")

print("Version: ", version)
print("Eager mode: ", executing_eagerly)
print("Hub Version: ", h_version)
print("GPU is", "available" if avai else "NOT AVAILABLE")

确保在您的GPU支持计算机中安装了最新的TensorFlow 2.x GPU，在python中执行以下代码，

from __future__ import absolute_import, division, print_function, unicode_literals

import tensorflow as tf 

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

会得到一个输出看起来像，

2020-02-07 10：45：37.587838：我tensorflow / stream_executor / cuda / cuda_gpu_executor.cc：1006]从SysFS读取成功的NUMA节点具有负值（-1），但必须至少有一个NUMA节点，因此返回NUMA节点为零2020-02-07 10：45：37.588896：I tensorflow / core / common_runtime / gpu / gpu_device.cc：1746]添加可见的gpu设备：0、1、2、3、4、5、6、7 Num可用GPU：8