tensorflow獲取可用GPU裝置
阿新 • • 發佈:2019-01-07
主要內容:
- 使用tensorflow查詢機器上是否存在可用的gpu裝置
- 使用tensorflow獲取可用的gpu裝置編號
- tensorflow對GPU裝置的編碼
使用tensorflow查詢機器上是否存在可用的gpu裝置
def is_gpu_available(cuda_only=True):
"""
code from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/platform/test.py
Returns whether TensorFlow can access a GPU.
Args:
cuda_only: limit the search to CUDA gpus.
Returns:
True iff a gpu device of the requested kind is available.
"""
from tensorflow.python.client import device_lib as _device_lib
if cuda_only:
return any((x.device_type == 'GPU')
for x in _device_lib.list_local_devices())
else:
return any((x.device_type == 'GPU' or x.device_type == 'SYCL')
for x in _device_lib.list_local_devices())
使用tensorflow獲取可用的gpu裝置編號
def get_available_gpus():
"""
code from http://stackoverflow.com/questions/38559755/how-to-get-current-available-gpus-in-tensorflow
"""
from tensorflow.python.client import device_lib as _device_lib
local_device_protos = _device_lib.list_local_devices()
return [x.name for x in local_device_protos if x.device_type == 'GPU']
tensorflow對GPU裝置的編碼
執行:
CUDA_VISIBLE_DEVICES=1,2 python test_util_tf.py
輸出為:
/gpu:0
/gpu:1
可以看出, 無論CUDA可見的裝置是哪幾個, tensorflow都會對它們從0開始重新編碼。