tensorflow获取可用GPU设备

正常来说,运行下面两句会打印tensorflow能用的CPU和GPU,以及相关的设备参数

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

以下罗列几个常用的查看service配置的函数:

使用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/False
    True if 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']

返回格式为:['/device:GPU:0', '/device:GPU:1']

---------------------------------------------------------------------------------------------------------------------------------------

参考博客:https://blog.csdn.net/littlehaes/article/details/82317220

                    https://blog.csdn.net/weixin_35653315/article/details/71403386

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值