TensorFlow-gpu 安装检测cuda版本程序

安装TensorFlow-gpu最大的问题就是显卡驱动版本与软件包中不匹配。这个python程序可以检测目前系统中的显卡驱动。由于条件限制,目前只能支持Windows系统。希望可以帮助到大家。

# -*- coding: utf-8 -*-

import os
import re
import subprocess
import io
import platform

if platform.system() != "Windows":
    print("目前只支持windows平台")
    exit(-1)

driver_runtime=[
    [441.22,'10.2.89'],
    [418.96,'10.1.105'],
    [411.31,'10.0.130'],
    [398.26,'9.2.148'],
    [397.44,'9.2.88'],
    [391.29,'9.1.85'],
    [385.54,'9.0.76'],
    [376.51,'8.0.61'],
    [369.30,'8.0.44'],
    [353.66,'7.5.16'],
    [347.62,'7.0.28']
    ]
'''
CUDA Toolkit	Linux x86_64 Driver Version	Windows x86_64 Driver Version
CUDA 10.2.89	>= 440.33	>= 441.22
CUDA 10.1 (10.1.105 general release, and updates)	>= 418.39	>= 418.96
CUDA 10.0.130	>= 410.48	>= 411.31
CUDA 9.2 (9.2.148 Update 1)	>= 396.37	>= 398.26
CUDA 9.2 (9.2.88)	>= 396.26	>= 397.44
CUDA 9.1 (9.1.85)	>= 390.46	>= 391.29
CUDA 9.0 (9.0.76)	>= 384.81	>= 385.54
CUDA 8.0 (8.0.61 GA2)	>= 375.26	>= 376.51
CUDA 8.0 (8.0.44)	>= 367.48	>= 369.30
CUDA 7.5 (7.5.16)	>= 352.31	>= 353.66
CUDA 7.0 (7.0.28)	>= 346.46	>= 347.62
'''
def compare_driver_toolkit(driver,toolkit):
    #driver 是电脑上显卡的驱动
    #toolkit 是Tensorflow-gpu安装的版本
    
    b_find=0
    for item in driver_runtime:
        if float(driver) > item[0]:
            b_find=1
            break
    if b_find:
        print("目前显卡的驱动:",item[1])
        print("Toolkit的版本:",toolkit)
        t=list(map(lambda x:int(x),list(toolkit.split("."))))
        d=list(map(lambda x:int(x),list(item[1].split("."))))

        a=d[0]*10000000+d[1]*100000+d[2]
        b=t[0]*10000000+t[1]*100000+t[2]
        return a > b
    else:
        return False

def execCmd(cmd):
    #popen会返回执行的命令
    return(os.popen(cmd).read())

def splitResult(result):
    result=' '.join(result.split())
    result=result.split(' ')
    if len(result)>2:
        return result

cmd="conda list"
#result=subprocess.check_output("dir").decode('utf-8')
result=execCmd(cmd)
result=result.split('\n')

i=0
for line in result:
    i=i+1
    idx=line.find('cudatoolkit')
    if idx>-1:
        cudatool=splitResult(line)
    idx=line.find('cudnn')
    if idx>-1:
        cudnn=splitResult(line)

print('===== cuda 安装信息:=====')
if len(cudatool)<2 and len(cudnn)<2:
    #
    print("没有正确安装cudatoolkit或者cudnn,请使用conda install cudatoolkit==版本 cudnn 安装")
else:
    print('cuda tool kit runtime 版本:',cudatool[1])
    print('cudnn 版本:',cudnn[1])
print("\n===== nvidia 显卡信息 =====")
cmd='nvidia-smi -q -u 2>&1'
result=execCmd(cmd)
#print(result)
idx=result.find('不是内部或外部命令')
if  idx>= 0:
    print('没有安装NVIDIA显卡!')
    exit(0)
else:
    result=result.split('\n')
    for line in result:
        idx = line.find('Driver Version')
        if idx >-1:
            driver_version=line.split(':')[1]
        idx = line.find('CUDA Version')
        if idx > -1:
            cuda_version=line.split(":")[1]
    print("Driver Version:",driver_version)
    print("CUDA Version:",cuda_version)

#目前系统中状态
print("\n===== 相关结果 =====")
if (compare_driver_toolkit(driver_version,cudatool[1])):
    print("显卡驱动与TensorFlow-gpu安装cudatoolkit匹配")
else:
    print("显卡驱动版本过低,请升级显卡驱动")
    print("显卡信息:",execCmd("nvidia-smi -L"))
    print("请去官方网址 https://www.nvidia.cn/Download/index.aspx?lang=cn 更新显卡驱动")

你可能感兴趣的