tensorflow获取GPU设备

转自：http://blog.csdn.net/weixin_35653315/article/details/71403386

主要内容：

• 使用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())
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

使用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']
1
2
3
4
5
6
7

完整代码见 tf.py, 测试代码见test_util_tf.py.

tensorflow对GPU设备的编码

执行：

CUDA_VISIBLE_DEVICES=1,2  python test_util_tf.py
1

输出为：

/gpu:0/gpu:1
1
2

可以看出， 无论CUDA可见的设备是哪几个， tensorflow都会对它们从0开始重新编码。