如何在tensorflow中获取当前可用的GPU?

时间:2016-07-25 04:30:38

标签: python gpu tensorflow

我计划使用分布式TensorFlow,我看到TensorFlow可以使用GPU进行培训和测试。在群集环境中,每台计算机可能有0个或1个或更多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']。我该如何实现呢?

12 个答案:

答案 0 :(得分:184)

有一个名为device_lib.list_local_devices()的未记录方法,可让您列出本地进程中可用的设备。 ( N.B。作为未记录的方法,这会受到向后不兼容的更改。)该函数返回DeviceAttributes protocol buffer个对象的列表。您可以按如下方式提取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()将运行一些初始化代码,默认情况下,它将分配所有设备(GitHub issue)上的所有GPU内存。要避免这种情况,请首先创建一个具有明确小per_process_gpu_fractionallow_growth=True的会话,以防止分配所有内存。有关详细信息,请参阅this question

答案 1 :(得分:90)

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

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

答案 2 :(得分:17)

test util中还有一个方法。 因此,所有要做的就是:

tf.test.is_gpu_available()

和/或

tf.test.gpu_device_name()

在Tensorflow文档中查找参数。

答案 3 :(得分:8)

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

因为目前只有Nvidia的gpus适用于NN框架,所以答案仅涵盖它们。 Nvidia has a page他们记录了如何使用/ proc文件系统界面获取有关驱动程序,任何已安装的NVIDIA显卡和AGP状态的运行时信息。

  

/proc/driver/nvidia/gpus/0..N/information

     

提供有关的信息   每个安装的NVIDIA图形适配器(型号名称,IRQ,BIOS   版本,总线类型)。请注意,BIOS版本仅适用于   X正在运行。

因此,您可以从命令行cat /proc/driver/nvidia/gpus/0/information运行此命令并查看有关您的第一个GPU的信息。它很容易run this from python,你也可以检查第二,第三,第四GPU直到它失败。

Mr Mrry的答案肯定更强大,我不确定我的答案是否适用于非Linux机器,但Nvidia的页面提供了其他有趣的信息,而这些信息并不是很多人所知。

答案 4 :(得分:8)

accepted answer为您提供了GPU的数量,但它也分配了这些GPU上的所有内存。您可以通过在调用device_lib.list_local_devices()之前创建一个内存较低的固定会话来避免这种情况,这对于某些应用程序可能是不需要的。

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

import subprocess

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

答案 5 :(得分:3)

在TensorFlow 2.0中,您可以使用&& sentText.react.me === false

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

如果您安装了两个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)

您可以在this documentation page上了解更多信息。

答案 6 :(得分:2)

以下在tensorflow 2.0中起作用:

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

答案 7 :(得分:2)

我的机器上装有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

答案 8 :(得分:1)

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

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

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

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

答案 9 :(得分:0)

使用这种方法并检查所有零件:

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

答案 10 :(得分:0)

确保在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:我 tensorflow / core / common_runtime / gpu / gpu_device.cc:1746]添加可见 gpu设备:0、1、2、3、4、5、6、7个可用GPU数量:8

答案 11 :(得分:0)

我正在研究TF-2.1和手电筒,所以我不想在任何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"\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%]

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