温馨提示×

Ubuntu上PyTorch与CUDA兼容性

小樊
58
2025-10-25 06:00:49
栏目: 智能运维

Ubuntu上PyTorch与CUDA兼容性指南
在Ubuntu系统上配置PyTorch与CUDA的兼容环境,需重点关注Ubuntu版本选择CUDA与PyTorch版本匹配驱动安装环境验证四大环节,以下是具体说明:

一、Ubuntu版本选择建议

Ubuntu的LTS(长期支持)版本对PyTorch及CUDA的支持更稳定,推荐优先选择:

  • Ubuntu 22.04 LTS (Jammy Jellyfish):默认软件包较新(如Python 3.10、GCC 11),支持CUDA 12.x及PyTorch 2.0+,适合新硬件(如RTX 40系列)和前沿研究;
  • Ubuntu 20.04 LTS (Focal Fossa):社区资源丰富,CUDA 11.x及以下版本支持成熟,适合工业级部署或依赖旧版框架(如TensorFlow 1.x)的场景。
    非LTS版本(如23.10)更新频繁、支持周期短(仅9个月),易导致依赖断裂,不建议用于生产环境。

二、PyTorch与CUDA版本兼容性表

PyTorch版本需与已安装的CUDA Toolkit版本严格匹配,以下是常见组合(来源:PyTorch官方文档及社区实践):

PyTorch版本 兼容CUDA版本 推荐cuDNN版本 适用场景
2.1.0+ 11.8 8.6 最新功能、RTX 40系列
2.0.1+ 11.7 8.5 新项目、高性能计算
1.13.1+ 11.6 8.3.2 视频生成、扩散模型
1.11.0+ 11.3 8.3 传统项目、工业部署
1.10.0+ 11.3/10.2 8.2 旧硬件、兼容性需求高

注:若使用PyTorch 1.10及以下版本,可选择CUDA 11.3或10.2;若使用PyTorch 2.0及以上版本,建议选择CUDA 11.7及以上。

三、安装步骤与注意事项

1. 安装NVIDIA驱动

驱动是CUDA运行的基础,需确保驱动版本与CUDA Toolkit兼容(可通过nvidia-smi查看当前驱动版本)。推荐通过以下命令安装:

# 添加官方驱动PPA
sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update
# 安装推荐驱动(如555,需根据显卡型号调整)
sudo apt install nvidia-driver-555
# 重启系统使驱动生效
sudo reboot

验证驱动:nvidia-smi(显示GPU信息及驱动版本即为成功)。

2. 安装CUDA Toolkit

  • 推荐方式(通过apt安装)
    以CUDA 11.3为例,执行以下命令:

    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb
    sudo dpkg -i cuda-keyring_1.0-1_all.deb
    sudo apt update
    sudo apt install -y cuda-11-3
    

    安装完成后,通过nvcc --version验证CUDA版本。

  • 手动安装(下载.run文件)
    从NVIDIA官网下载对应版本的CUDA Toolkit(如CUDA 11.7),运行安装脚本并遵循提示完成安装。

3. 安装PyTorch

  • 推荐方式(通过conda安装,自动管理依赖)
    以PyTorch 1.11.0 + CUDA 11.3为例:
    conda create -n pytorch_env python=3.8
    conda activate pytorch_env
    conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch
    
  • 手动安装(通过pip安装)
    以PyTorch 1.11.0 + CUDA 11.3为例:
    pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113
    

4. 配置环境变量(可选)

若安装CUDA Toolkit时未自动配置,需手动添加以下环境变量至~/.bashrc

export PATH=/usr/local/cuda-11.3/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.3/lib64:$LD_LIBRARY_PATH

执行source ~/.bashrc使更改生效。

四、验证安装

安装完成后,通过以下Python代码验证PyTorch是否能识别CUDA:

import torch
print("PyTorch版本:", torch.__version__)
print("CUDA可用性:", torch.cuda.is_available())  # 应返回True
print("CUDA版本:", torch.version.cuda)          # 应与安装的CUDA版本一致
print("GPU数量:", torch.cuda.device_count())    # 显示GPU数量
print("GPU名称:", torch.cuda.get_device_name(0)) # 显示首个GPU名称

torch.cuda.is_available()返回True,则表示环境配置成功。

常见问题排查

  • torch.cuda.is_available()返回False
    1. 检查NVIDIA驱动是否正确安装(nvidia-smi是否有输出);
    2. 确认CUDA Toolkit版本与PyTorch版本匹配;
    3. 确保环境变量(PATHLD_LIBRARY_PATH)配置正确。
  • 非法内存访问错误
    检查CUDA架构是否匹配(如Ampere架构需设置TORCH_CUDA_ARCH_LIST=8.6),或升级PyTorch至最新版本。

0