PyTorch是Meta AI（Facebook）开发的开源机器学习框架，基于Torch库，用于计算机视觉和自然语言处理等应用。正确选择PyTorch版本与CUDA版本的兼容性是深度学习环境配置中的关键环节，直接影响GPU加速性能和模型训练效率。

CUDA是NVIDIA开发的GPU加速库，仅能在NVIDIA系列显卡中使用，AMD及Intel等公司出品的显卡不在本系列教程的考虑范围内。

前置教程

如想快速开始本教程，你可能需要先完成以下前置教程：

1. 版本选择的重要性

PyTorch与CUDA版本的兼容性问题常常导致以下情况：

  • GPU无法被PyTorch识别
  • 模型训练时出现CUDA相关错误
  • 性能无法充分利用GPU资源
  • 某些新功能不可用

因此，理解版本兼容性规则并选择合适的PyTorch版本至关重要。

2. CUDA版本与PyTorch兼容性

2.1 CUDA版本查看

在安装PyTorch之前，首先需要确认系统的CUDA版本。在WSL2中运行以下命令：

# 查看驱动支持的CUDA版本（仅供参考）
nvidia-smi

# 查看实际安装的CUDA Toolkit版本（以此为准）
nvcc --version

查看实际安装的CUDA Toolkit版本（以此为准）示意图

nvidia-smi 显示的是驱动支持的最高CUDA版本，而 nvcc --version 显示的是实际安装的CUDA Toolkit版本。应以 nvcc --version 为准来选择PyTorch版本。

如果你的系统中尚未安装CUDA Toolkit，即运行nvcc --version后没有任何内容显示或报错，没有关系，PyTorch 安装包自带 CUDA 运行时库，不需要系统级 CUDA Toolkit，直接安装PyTorch即可。但需确认你的显卡驱动支持该版本的CUDA Toolkit，你可以前往这个页面查看具体匹配关系：pytorch/RELEASE.md

3. PyTorch安装方法

3.1 环境准备

conda create -n PyTorch_env python=3.10 -y
conda activate PyTorch_env

3.2 安装PyTorch

根据上方的兼容性表格选择对应的版本，这里以本机CUDA 12.0为例，安装PyTorch 2.3：

# CUDA 12.0 + PyTorch 2.3
pip install torch==2.3 torchvision torchaudio --index-url https://download.PyTorch.org/whl/cu120 -i https://pypi.tuna.tsinghua.edu.cn/simple
  • torch==2.3-指定PyTorch版本
  • torchvision 和 torchaudio 是PyTorch的子库，用于计算机视觉和音频处理
  • --index-url 指定PyTorch官方CUDA版本的下载索引，确保安装正确CUDA版本的预编译包
  • -i 指定Python包索引源，加速安装速度
  • cu120 指定CUDA版本，确保PyTorch与CUDA兼容

4. 环境配置与验证

4.1 验证CUDA支持

# 完整验证脚本
python -c "
import torch
print(f'PyTorch version: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'CUDA version: {torch.version.cuda}')
    print(f'Device count: {torch.cuda.device_count()}')
    print(f'Current device: {torch.cuda.current_device()}')
    print(f'Device name: {torch.cuda.get_device_name(0)}')
    print(f'Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB')
else:
    print('CUDA不可用，请检查安装')
"

完整验证脚本示意图

  • PyTorch version: 2.3.0+cu121：PyTorch版本
  • CUDA available: True：CUDA可用
  • CUDA version: 12.0：CUDA版本
  • Device count: 1：GPU数量
  • Current device: 0：当前GPU设备
  • Device name: GeForce RTX 2080：GPU名称
  • Memory: 22 GB：GPU内存

4.2 测试GPU计算

# 简单GPU计算测试
python -c "
import torch
if torch.cuda.is_available():
    x = torch.randn(1000, 1000).cuda()
    y = torch.randn(1000, 1000).cuda()
    z = torch.matmul(x, y)
    print('[OK] GPU计算测试成功')
    print(f'  结果张量设备: {z.device}')
    print(f'  结果形状: {z.shape}')
else:
    print('[失败] CUDA不可用，请检查安装')
"

简单GPU计算测试示意图

4.3 性能基准测试

# 快速基准测试
python -c "
import torch
import time

if torch.cuda.is_available():
    # CPU测试
    cpu_tensor = torch.randn(3000, 3000)
    start = time.time()
    cpu_result = cpu_tensor @ cpu_tensor
    cpu_time = time.time() - start

    # GPU测试
    gpu_tensor = cpu_tensor.cuda()
    torch.cuda.synchronize()
    start = time.time()
    gpu_result = gpu_tensor @ gpu_tensor
    torch.cuda.synchronize()
    gpu_time = time.time() - start

    print(f'CPU时间: {cpu_time:.4f}s')
    print(f'GPU时间: {gpu_time:.4f}s')
    print(f'加速比: {cpu_time/gpu_time:.2f}x')
"

快速基准测试示意图