有一个名为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创建一个会话，以防止分配所有内存。请参阅这个问题了解更多细节。

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

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

除了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/ GPU /0/information运行这个命令，并查看关于你的第一个GPU的信息。从python中运行这个很容易，你也可以检查第二个、第三个、第四个GPU，直到它失败。

当然，Mrry的答案更加可靠，我不确定我的答案是否适用于非linux机器，但Nvidia的页面提供了其他有趣的信息，这些信息不是很多人知道的。

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

tf.test.is_gpu_available()

和/或

tf.test.gpu_device_name()

在Tensorflow文档中查找参数。

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

我最终使用nvidia-smi来获得gpu的数量，而不分配任何内存。

import subprocess

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

从TensorFlow 2.1开始，你可以使用tf.config.list_physical_devices('GPU'):

import tensorflow as tf

gpus = tf.config.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

在TF 2.0中，您必须添加experimental:

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

See:

引导页 当前的API

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开始，你可以放弃实验性:

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

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

用这种方法检查所有部件:

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")

确保你有最新的TensorFlow 2。x GPU安装在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。答:1006]成功 从SysFS读取的NUMA节点有负值(-1)，但必须有 至少一个NUMA节点，因此返回NUMA节点零2020-02-07 10:45:37.588896:我 tensorflow /核心/ common_runtime / gpu / gpu_device。cc:1746]增加可见 gpu设备:0,1,2,3,4,5,6,7 Num gpu可用:8

我在我的机器上有一个名为NVIDIA GTX GeForce 1650 Ti的GPU, tensorflow-gpu==2.2.0

运行以下两行代码:

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

输出:

Num GPUs Available:  1

我正在TF-2.1和torch上工作，所以我不想在任何ML框架中指定这个自动选择。我只使用原版的nvidia-smi和os。找到一个空的显卡。

def auto_gpu_selection(usage_max=0.01, mem_max=0.05):
"""Auto set CUDA_VISIBLE_DEVICES for gpu

:param mem_max: max percentage of GPU utility
:param usage_max: max percentage of GPU memory
:return:
"""
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
log = str(subprocess.check_output("nvidia-smi", shell=True)).split(r"\n")[6:-1]
gpu = 0

# Maximum of GPUS, 8 is enough for most
for i in range(8):
    idx = i*3 + 2
    if idx > log.__len__()-1:
        break
    inf = log[idx].split("|")
    if inf.__len__() < 3:
        break
    usage = int(inf[3].split("%")[0].strip())
    mem_now = int(str(inf[2].split("/")[0]).strip()[:-3])
    mem_all = int(str(inf[2].split("/")[1]).strip()[:-3])
    # print("GPU-%d : Usage:[%d%%]" % (gpu, usage))
    if usage < 100*usage_max and mem_now < mem_max*mem_all:
        os.environ["CUDA_VISIBLE_EVICES"] = str(gpu)
        print("\nAuto choosing vacant GPU-%d : Memory:[%dMiB/%dMiB] , GPU-Util:[%d%%]\n" %
              (gpu, mem_now, mem_all, usage))
        return
    print("GPU-%d is busy: Memory:[%dMiB/%dMiB] , GPU-Util:[%d%%]" %
          (gpu, mem_now, mem_all, usage))
    gpu += 1
print("\nNo vacant GPU, use CPU instead\n")
os.environ["CUDA_VISIBLE_EVICES"] = "-1"

如果我能得到任何GPU，它将CUDA_VISIBLE_EVICES设置为该GPU的BUSID:

GPU-0 is busy: Memory:[5738MiB/11019MiB] , GPU-Util:[60%]
GPU-1 is busy: Memory:[9688MiB/11019MiB] , GPU-Util:[78%]

Auto choosing vacant GPU-2 : Memory:[1MiB/11019MiB] , GPU-Util:[0%]

else，设置为-1使用CPU:

GPU-0 is busy: Memory:[8900MiB/11019MiB] , GPU-Util:[95%]
GPU-1 is busy: Memory:[4674MiB/11019MiB] , GPU-Util:[35%]
GPU-2 is busy: Memory:[9784MiB/11016MiB] , GPU-Util:[74%]

No vacant GPU, use CPU instead

注意:在导入任何需要GPU的ML帧之前使用这个函数，然后它会自动选择一个GPU。此外，你可以轻松设置多个任务。

在TensorFlow Core v2.3.0中，以下代码应该可以工作。

import tensorflow as tf
visible_devices = tf.config.get_visible_devices()
for devices in visible_devices:
  print(devices)

根据您的环境，这段代码将产生流动的结果。

PhysicalDevice (name = / physical_device: CPU: 0, device_type = CPU) PhysicalDevice (name = / physical_device: GPU: 0, device_type = GPU)

tensorflow推荐的最新版本:

tf.config.list_physical_devices('GPU')

在任何shell中运行以下命令

python -c "import tensorflow as tf; print(\"Num GPUs Available: \", len(tf.config.list_physical_devices('GPU')))"

您可以使用以下代码字段来显示设备名称、类型、内存和位置。

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())