[python]python onnxruntime gpu是否可用

在Python中,要检查ONNX Runtime是否可以使用GPU,可以使用以下代码:

import onnxruntime as ort
 
# 检查CUDA是否可用
def is_cuda_available():
    try:
        return ort.get_device() == 'GPU'
    except Exception:
        return False
 
# 检查ONNX Runtime是否支持CUDA
def is_onnxruntime_cuda_supported():
    return ort.get_device() == 'GPU'
 
# 检查ONNX Runtime是否使用CUDA
def is_onnxruntime_using_cuda():
    providers = ort.get_available_providers()
    return 'CUDA' in providers and 'CUDAExecutionProvider' in providers
 
# 测试代码
print(f"CUDA Available: {is_cuda_available()}")
print(f"ONNX Runtime CUDA Supported: {is_onnxruntime_cuda_supported()}")
print(f"ONNX Runtime Using CUDA: {is_onnxruntime_using_cuda()}")

请确保已经安装了ONNX Runtime,并且如果使用GPU版本,确保已经安装了相应的CUDA驱动和cuDNN库。如果系统中没有可用的GPU,则这些函数将返回False

你可能感兴趣的:(Python,python,人工智能,深度学习)