Ubuntu 上安装 PyTorch 与 CUDA 的简明步骤
一 准备与检查
- 更新系统并安装基础工具:sudo apt update && sudo apt upgrade -y;sudo apt install -y build-essential cmake git wget unzip python3 python3-pip
- 确认 GPU 与驱动:lspci | grep -i nvidia;nvidia-smi(右上角显示 Supported/Runtime CUDA,表示驱动支持的最高 CUDA 运行时)
- 创建隔离环境(推荐):使用 Miniconda/Anaconda 或 venv,便于多版本管理
- conda create -n pytorch-env python=3.10 && conda activate pytorch-env
- python3 -m venv pytorch_env && source pytorch_env/bin/activate
- 说明:如无 NVIDIA GPU,可直接安装 CPU 版 PyTorch;如有 GPU,优先使用 PyTorch 官方预编译包(自带 CUDA 运行时),避免与系统 CUDA 版本强绑定。
二 两种安装方案
- 方案 A(推荐)使用官方预编译包(自带 CUDA 运行时)
- 在激活的虚拟环境中,按官网生成命令安装,例如 CUDA 11.8:
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
- 其他常见版本示例:CUDA 11.3:pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
- CPU 版:pip install torch torchvision torchaudio
- 说明:此方式无需单独安装系统级 CUDA/cuDNN,PyTorch 二进制自带所需运行时,最省事。
- 方案 B 系统级安装 CUDA 与 cuDNN(适合需要系统库的场景)
- 安装与驱动匹配的 CUDA(示例 Ubuntu 22.04 + CUDA 11.8):
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget http://developer.download.nvidia.com/compute/cuda/11.8.1/local_installers/cuda-repo-ubuntu2204-11-8-local-11.8.1-465.19.06-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-11-8-local-11.8.1-465.19.06-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2204-11-8-local-11.8.1-465.19.06-1/7fa2af80.pub
sudo apt update && sudo apt install -y cuda
- 安装与 CUDA 匹配的 cuDNN(示例 11.8):
tar -xzvf cudnn-11.8.1-linux-x64-v8.6.1.32.tgz
sudo cp cuda/include/cudnn*.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*
- 设置环境变量(若 CUDA 安装到 /usr/local/cuda):
echo ‘export PATH=/usr/local/cuda/bin:$PATH’ >> ~/.bashrc
echo ‘export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH’ >> ~/.bashrc
source ~/.bashrc
- 说明:系统级安装更灵活,但依赖管理更复杂;务必保证 CUDA 与 cuDNN 版本匹配。
三 验证安装
- 基础验证:
python - <<‘PY’
import torch, sys
print(“PyTorch version:”, torch.version)
print(“CUDA available:”, torch.cuda.is_available())
if torch.cuda.is_available():
print(“CUDA version:”, torch.version.cuda)
print(“Device count:”, torch.cuda.device_count())
print(“Current device:”, torch.cuda.current_device())
print(“Device name:”, torch.cuda.get_device_name(0))
PY
- 期望输出要点:torch.cuda.is_available() 为 True;torch.version.cuda 显示 PyTorch 二进制内的 CUDA 运行时版本(与 nvidia-smi 右上角的“Supported/Runtime CUDA”概念不同,二者不必完全一致)。
四 常见问题与排错
- 版本不匹配与概念澄清:nvidia-smi 右上角显示的是驱动支持的“最高 CUDA 运行时”,而 PyTorch 的 torch.version.cuda 是其自带运行时的版本;优先以 PyTorch 官网命令与二进制为准,避免强行对齐二者版本号。
- 驱动问题:若 nvidia-smi 无法运行或报错,先正确安装/更新 NVIDIA 驱动(可参考“附加驱动”或手动安装方式),再安装 CUDA/cuDNN。
- 环境变量未生效:若采用系统级 CUDA,确认 ~/.bashrc 中 PATH 与 LD_LIBRARY_PATH 已设置并 source 生效;执行 which nvcc 与 nvcc -V 检查工具链是否可用。
- 多版本 CUDA 切换:可通过调整 PATH/LD_LIBRARY_PATH 指向不同 /usr/local/cuda-xx.x 实现切换;或直接使用方案 A 规避多版本管理复杂度。