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 规避多版本管理复杂度。