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

5 浏览
0 Comments

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

我计划使用分布式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，所以我不想从操作系统内核中获取GPU信息的方法。

简而言之，我希望有一个类似于tf.get_available_gpus()的函数，如果机器上有两个可用的GPU，它将返回['/gpu:0', '/gpu:1']。我该如何实现这个？

0