我有计划使用分布式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'].我该如何实现呢？