1. 程式人生 > >tensorflow獲取可用GPU裝置

tensorflow獲取可用GPU裝置

主要內容:

  • 使用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開始重新編碼。