如何在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,所以我不想知道从OS内核获取GPU信息的方法.

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

Answer 1

有一个未记录的方法device_lib.list_local_devices(),使您可以列出本地进程中可用的设备.(NB作为未记录的方法,这会受到向后不兼容的更改.)该函数返回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,以防止分配所有内存.有关详细信息,请参阅此问题.

Answer 2

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

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

Answer 3

测试工具中还有一个方法.所以必须做的就是:

tf.test.is_gpu_available()

和/或

tf.test.gpu_device_name()

查找Tensorflow文档以获取参数.

Answer 4

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

我最终使用nvidia-smi来获取GPU的数量而不分配任何内存.

import subprocess

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

Answer 5

在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

您可以在此文档页面上了解更多信息。

Answer 6

除了Mrry的优秀解释,他建议使用device_lib.list_local_devices()我可以告诉你如何从命令行检查GPU相关信息.

因为目前只有Nvidia的gpus适用于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的页面提供了其他有趣的信息,而这些信息并不是很多人所知.

Answer 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

Answer 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

Answer 9

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