如何在张量流中获取当前可用的GPU？

小编典典

如何在张量流中获取当前可用的GPU？

python

我有一个使用分布式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']。我该如何实施？

阅读 179

2020-12-20

共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']

请注意（至少在TensorFlow1.4之前），调用device_lib.list_local_devices()将运行一些初始化代码，默认情况下，这些初始化代码将在所有设备上分配所有GPU内存（GitHubissue）。为避免这种情况，请首先使用一个显着小的per_process_gpu_fraction或创建一个会话allow_growth=True，以防止分配所有内存。

2020-12-20