如何判断你的PyTorch是GPU版还是CPU版？

PyTorch作为当前最流行的深度学习框架之一，支持在CPU和GPU(NVIDIA CUDA)上运行。对于深度学习开发者来说，正确识别PyTorch版本至关重要，因为GPU版本可以带来10-100倍的性能提升。本文将全面介绍如何判断你的PyTorch安装版本，并提供详细的案例分析和问题解决方案。

为什么需要区分GPU和CPU版本？

性能差异

GPU版本的PyTorch可以利用NVIDIA显卡的CUDA核心进行并行计算：

训练速度通常比CPU快10-100倍
能够处理更大的batch size
支持更复杂的模型架构

硬件要求

GPU版本需要满足以下条件：

兼容的NVIDIA显卡（如RTX 30/40系列、Tesla系列等）
正确安装的NVIDIA驱动和CUDA工具包
与硬件匹配的PyTorch GPU版本

如何检查PyTorch版本？

方法1：使用命令行快速检查

运行以下命令获取基本信息：

bash 复制代码

python -c "import torch; print(torch.__version__); print('CUDA available:', torch.cuda.is_available()); print('Device count:', torch.cuda.device_count())"

输出案例1：GPU版本正常工作

复制代码

2.3.0+cu121
CUDA available: True
Device count: 1

解读：

+cu121表示PyTorch编译时使用的CUDA版本是12.1
CUDA available: True表示CUDA可用
Device count: 1表示检测到1块可用GPU

输出案例2：CPU版本

复制代码

2.3.0
CUDA available: False
Device count: 0

解读：

版本号没有+cuxx后缀，表示是CPU版本
CUDA available: False确认不支持CUDA

方法2：使用详细检查脚本

python 复制代码

import torch

print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"CUDA版本: {torch.version.cuda}")
    print(f"GPU设备数量: {torch.cuda.device_count()}")
    print(f"当前设备: {torch.cuda.current_device()}")
    print(f"设备名称: {torch.cuda.get_device_name(0)}")
    print(f"设备内存: {torch.cuda.get_device_properties(0).total_memory/1024**3:.2f} GB")
else:
    print("当前安装的是CPU版PyTorch或CUDA不可用")

输出案例：详细GPU信息

复制代码

PyTorch版本: 2.3.0+cu121
CUDA可用: True
CUDA版本: 12.1
GPU设备数量: 1
当前设备: 0
设备名称: NVIDIA GeForce RTX 4090
设备内存: 24.00 GB

常见问题与解决方案

问题1：安装了GPU版但显示不可用

可能原因：

NVIDIA驱动未正确安装
CUDA工具包版本不匹配
PyTorch版本与CUDA版本不兼容

解决方案：

检查NVIDIA驱动：运行nvidia-smi
检查CUDA版本：nvcc --version
重新安装匹配版本的PyTorch