如何在PyTorch中获取GPU的CUDA计算能力？解决兼容检测需求

用PyTorch代码直接获取计算能力 #

不用依赖外部工具，直接通过PyTorch的API就能拿到GPU的CUDA计算能力——哪怕当前PyTorch版本不支持该GPU，只要CUDA驱动正常安装，就能执行这段代码提前判断：

import torch
if torch.cuda.is_available():
    # 获取指定GPU的计算能力，返回元组(主版本号, 次版本号)
    cap = torch.cuda.get_device_capability(0)
    cap_version = cap[0] + cap[1]/10
    print(f"当前GPU的CUDA计算能力: {cap_version}")
    # 对比PyTorch要求的最低版本
    if cap_version < 3.7:
        print("该GPU不兼容当前PyTorch版本，请降级PyTorch或更换GPU")
else:
    print("未检测到可用的CUDA设备")

终端命令直接查询 #

如果不想写代码，也可以用NVIDIA自带的命令行工具查询：

nvidia-smi --query-gpu=compute_cap --format=csv,noheader

执行后会直接输出GPU的CUDA计算能力数值，比如3.5。

补充说明 #

torch.cuda.is_available()仅能判断系统是否识别到CUDA设备、驱动是否正常，无法校验PyTorch对该GPU计算能力的支持性。所以必须额外检查计算能力，避免运行时出现不兼容报错。

如果你的GPU计算能力低于3.7，有两个解决方向：

• 降级PyTorch到支持CUDA 3.5的旧版本；
• 更换符合计算能力要求的GPU。

内容的提问来源于stack exchange，提问作者pete