tf_utils.py 文件源码

def has_gpu() -> bool:
    """Check if TensorFlow can access GPU.

    The test is based on
        https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/platform/test.py
    ...but we are interested only in CUDA GPU devices.

    Returns:
        True, if TF can access the GPU
    """
    # pylint: disable=global-statement
    global __HAS_GPU_RESULT
    # pylint: enable=global-statement
    if __HAS_GPU_RESULT is None:
        __HAS_GPU_RESULT = any((x.device_type == 'GPU')
                               for x in _device_lib.list_local_devices())
    return __HAS_GPU_RESULT