如何在张量流中获取当前可用的GPU?
python
tensorflow
5
0

我有一个使用分布式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信息的方法。

简而言之,我想要一个类似tf.get_available_gpus()的函数['/gpu:0', '/gpu:1']如果机器中有两个GPU,它将返回['/gpu:0', '/gpu:1'] 。我该如何实施?

参考资料:
Stack Overflow
收藏
评论
共 6 个回答
高赞 时间 活跃

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

tf.test.is_gpu_available()

和/或

tf.test.gpu_device_name()

在Tensorflow文档中查找参数。

收藏
评论

有一个未device_lib.list_local_devices()方法,称为device_lib.list_local_devices() ,使您可以列出本地进程中可用的设备。 ( 注意,作为一个未DeviceAttributes方法,此方法可能会向后不兼容更改。)该函数返回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_fractionallow_growth=True创建会话,以防止分配所有内存。有关更多详细信息,请参见此问题

收藏
评论

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

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

import subprocess

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

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

看到:

收藏
评论

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

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/gpus/0/information运行此命令,并查看有关第一个GPU的信息。 从python运行它很容易,并且您可以检查第二,第三,第四GPU直到失败。

肯定地Mrry的答案更可靠,而且我不确定我的答案是否可以在非Linux机器上使用,但是Nvidia的页面提供了其他有趣的信息,但鲜为人知。

收藏
评论
新手导航
  • 社区规范
  • 提出问题
  • 进行投票
  • 个人资料
  • 优化问题
  • 回答问题

关于我们

常见问题

内容许可

联系我们

@2020 AskGo
京ICP备20001863号