使用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']

 

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

tensorflow对GPU设备的编码

执行：

CUDA_VISIBLE_DEVICES=1,2  python test_util_tf.py

 

输出为：

/gpu:0
/gpu:1

 

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

 

Tensorflow 使用时cpu编译不支持警告

 原因是下载TensorFlow的版本不支持cpu的AVX2编译。

可能是因为安装时使用的pip install tensorflow ，这样默认会下载X86_64的SIMD版本。

有两种解决办法：

1.忽略这个警告，不看它！

import os  
os.environ["TF_CPP_MIN_LOG_LEVEL"]='1' # 这是默认的显示等级，显示所有信息  
os.environ["TF_CPP_MIN_LOG_LEVEL"]='2' # 只显示 warning 和 Error   
os.environ["TF_CPP_MIN_LOG_LEVEL"]='3' # 只显示 Error

2.彻底解决，换成支持cpu用AVX2编译的TensorFlow版本。

首先，卸载原来版本的TensorFlow

pip uninstall tensorflow

完成后可以用

pip list

查看一下所有库，检查是否成功卸载。

然后去github下载正确的tf版本，Windows点这里这里下载。其他操作系统可以点这里这里找到对应的版本。博主用的win10，python3.6，所以选择如下图：

对应的path在上边的找到对应的.whl下载即可。

然后安装：

pip install tensorflow-1.6.0-cp36-cp36m-win_amd64.whl

成功安装后，再次运行TensorFlow就不会报错啦~