Python 代码查看本机 CUDA 版本完整教程博文

14次阅读

前言

做深度学习、AI 模型训练时，经常需要确认环境 CUDA 版本，用来匹配 PyTorch、TensorFlow、CUDA Toolkit。纯 Python 脚本无需命令行，直接运行就能输出 CUDA 版本，分两种场景：PyTorch 环境、原生 CUDA C 拓展检测，下面附上可直接复制运行的完整代码。

一、基于 PyTorch 获取 CUDA 版本（深度学习常用）

绝大多数炼丹环境都安装了 PyTorch，这套代码同时输出：CUDA 是否可用、系统 CUDA 驱动版本、PyTorch 编译绑定的 CUDA 版本、GPU 设备信息。

完整代码

python

运行

import torch

def get_cuda_info():
    print("===== CUDA 环境信息检测 =====")
    # 判断CUDA是否可用
    cuda_available = torch.cuda.is_available()
    print(f"CUDA 是否可用: {cuda_available}")

    if not cuda_available:
        print("未检测到可用NVIDIA显卡/CUDA环境")
        return

    # 系统驱动支持的最高CUDA版本
    driver_cuda_ver = torch.version.cuda
    print(f"显卡驱动支持CUDA版本: {driver_cuda_ver}")

    # PyTorch编译时绑定的CUDA版本
    torch_cuda_ver = torch.version.cuda
    print(f"PyTorch内置CUDA版本: {torch_cuda_ver}")

    # GPU数量与设备名
    gpu_count = torch.cuda.device_count()
    print(f"GPU 设备总数: {gpu_count}")
    for i in range(gpu_count):
        gpu_name = torch.cuda.get_device_name(i)
        print(f"GPU{i} 型号: {gpu_name}")

if __name__ == "__main__":
    get_cuda_info()

运行结果示例

plaintext

===== CUDA 环境信息检测 =====
CUDA 是否可用: True
显卡驱动支持CUDA版本: 12.1
PyTorch内置CUDA版本: 12.1
GPU 设备总数: 1
GPU0 型号: NVIDIA GeForce RTX 4060 Ti

说明

torch.version.cuda：返回当前驱动支持的 CUDA 大版本，比如 11.8、12.1
若返回None代表无 CUDA 环境，只能跑 CPU 训练

二、不依赖 PyTorch，原生 ctypes 读取 CUDA 版本

未安装深度学习框架，仅装了 CUDA Toolkit，用 Python 内置ctypes调用 cuda 动态库读取版本，跨 Windows/Linux 通用。

完整代码

python

运行

import ctypes
import os
import platform

def get_raw_cuda_version():
    # 适配不同系统cuda库文件名
    system = platform.system()
    if system == "Windows":
        lib_name = "nvcuda.dll"
    elif system == "Linux":
        lib_name = "libcuda.so"
    else:
        print("MacOS无原生CUDA支持")
        return None

    try:
        cuda_lib = ctypes.CDLL(lib_name)
    except OSError:
        print(f"未找到CUDA库 {lib_name}，未安装CUDA驱动")
        return None

    # 定义版本存储变量
    version = ctypes.c_int()
    # 调用cuDriverGetVersion API
    status = cuda_lib.cuDriverGetVersion(ctypes.byref(version))
    if status != 0:
        print("读取CUDA版本失败")
        return None

    ver_num = version.value
    major = ver_num // 1000
    minor = (ver_num % 1000) // 10
    full_version = f"{major}.{minor}"
    print(f"原生CUDA驱动版本: {full_version}")
    return full_version

if __name__ == "__main__":
    print("===== 原生CUDA驱动版本检测 =====")
    get_raw_cuda_version()

运行效果

plaintext

===== 原生CUDA驱动版本检测 =====
原生CUDA驱动版本: 12.2

三、常见问题解决

运行报错No module named torch 解决方案：pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
CUDA 显示 False，显卡识别不到
- 确认显卡是 NVIDIA，AMD/Intel 核显不支持 CUDA
- 更新 NVIDIA 显卡驱动
- 重装匹配版本 PyTorch
PyTorch CUDA 版本和驱动版本不一致驱动版本数字 ≥ PyTorch CUDA 版本即可向下兼容，例如驱动 12.2 可跑 cu118、cu121 的 PyTorch