如何获取 TensorFlow 中当前可用的 GPU？-IT科技

摘要：问题描述：我计划使用分布式 TensorFlow，并且我发现 TensorFlow 可以使用 GPU 进行训练和测试。在集群环境中，每台机器可以有 0 个或 1 个或更多 GPU，我希望在尽可能多的机器上的 GPU 上运行我的 TensorFlow 图表。我发现运行tf.Session()TensorFlow...

问题描述：

我计划使用分布式 TensorFlow，并且我发现 TensorFlow 可以使用 GPU 进行训练和测试。在集群环境中，每台机器可以有 0 个或 1 个或更多 GPU，我希望在尽可能多的机器上的 GPU 上运行我的 TensorFlow 图表。

我发现运行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']

解决方案 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，以防止分配所有内存。有关更多详细信息，请参阅此问题。

解决方案 2：

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

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

解决方案 3：

从 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')

看：

指南页面
当前 API

解决方案 4：

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

tf.test.is_gpu_available()

和/或

tf.test.gpu_device_name()

查阅 Tensorflow 文档以了解参数。

解决方案 5：

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

我最终使用 nvidia-smi 来获取 GPU 的数量，而无需在它们上分配任何内存。

import subprocess

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

解决方案 6：

除了 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 运行时可用。

因此，您可以从命令行运行它并查看有关第一个 GPU 的信息。从 Python 运行它cat /proc/driver/nvidia/gpus/0/information很容易，您还可以检查第二、第三、第四个 GPU，直到它失败。

毫无疑问，Mrry 的回答更为有力，而且我不确定我的回答是否能在非 Linux 机器上起作用，但 Nvidia 的页面提供了其他有趣的信息，但很多人都不知道。

解决方案 7：

以下在 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

解决方案 8：

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

解决方案 9：

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

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

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

物理设备（名称='/physical_device:CPU:0'，设备类型='CPU'）物理设备（名称='/physical_device:GPU:0'，设备类型='GPU'）

解决方案 10：

tensorflow推荐的最新版本：

tf.config.list_physical_devices('GPU')

解决方案 11：

我正在研究 TF-2.1 和 torch，所以我不想在任何 ML 框架中指定这种自动选择。我只是使用原始的nvidia-smi和os.environ来获得一个空闲的 gpu。

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"
")[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("
Auto choosing vacant GPU-%d : Memory:[%dMiB/%dMiB] , GPU-Util:[%d%%]
" %
              (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("
No vacant GPU, use CPU instead
")
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%]

否则，设置为-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。此外，您还可以轻松设置多个任务。

解决方案 12：

使用此方法并检查所有部件：

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

解决方案 13：

确保你的 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：I 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 可用 GPU 数量：8

解决方案 14：

在任何 shell 中运行以下命令

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

解决方案 15：

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

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

解决方案 16：

接受的答案为您提供了如下设备描述：

['/device:GPU:0']

如果您需要更多详细信息，可以使用tf.config.experimental.get_device_details()

import tensorflow as tf

def get_available_gpus():
        physical_gpus = tf.config.list_physical_devices(device_type="GPU")
        return [(x, tf.config.experimental.get_device_details(x)) for x in physical_gpus]

这将为您提供有关 device_name 和 compute_capability 的详细信息，例如：

[(PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'), {'device_name': 'NVIDIA T500', 'compute_capability': (7, 5)})]