在 Ubuntu 上配置 PyTorch GPU 支持
一 准备与检查
- 确认硬件与系统:需要 NVIDIA GPU 且支持 CUDA。在终端执行 nvidia-smi,若能看到 GPU 型号、驱动版本、CUDA 运行时版本,说明驱动与 GPU 识别正常。若命令不存在或报错,需先安装/修复 NVIDIA 驱动。建议优先使用系统推荐驱动,执行 ubuntu-drivers devices 查看推荐版本,并用 sudo ubuntu-drivers autoinstall 安装,随后重启。以上步骤可确保后续 CUDA 与 PyTorch 能正确调用 GPU。
二 安装驱动与 CUDA 工具链
- 安装驱动(两种常用方式,二选一)
- 推荐方式:使用系统推荐驱动
- 执行:sudo ubuntu-drivers devices → 查看推荐版本
- 执行:sudo ubuntu-drivers autoinstall
- 执行:sudo reboot
- 指定版本安装(示例为安装 470 系列)
- 执行:sudo apt-get update && sudo apt-get install nvidia-driver-470
- 执行:sudo reboot
- 安装 CUDA Toolkit(两种常用方式,二选一)
- 方式 A(Ubuntu 仓库):sudo apt-get update && sudo apt-get install cuda-11.0(示例版本,可按需调整)
- 方式 B(NVIDIA 官方安装器):前往 developer.nvidia.com/cuda-downloads 选择对应 Ubuntu 版本与 CUDA 版本下载 runfile,执行示例:
- wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda_11.7.0_515.43.04_linux.run
- sudo sh cuda_11.7.0_515.43.04_linux.run
- 配置环境变量(如使用官方安装器或需自定义路径)
- 编辑 ~/.bashrc:nano ~/.bashrc
- 在文件末尾添加(将 替换为实际版本,如 11.7):
- export PATH=/usr/local/cuda-/bin${PATH:+:${PATH}}
- export LD_LIBRARY_PATH=/usr/local/cuda-/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
- 使配置生效:source ~/.bashrc
- 安装 cuDNN(与 CUDA 版本匹配)
- 前往 developer.nvidia.com/cudnn 下载与 CUDA 兼容的版本,解压并按需复制到 CUDA 目录(示例命令,文件名按实际替换):
- tar -xzvf cudnn--linux-x64-v.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**
- 说明
- 若使用 conda 安装 PyTorch,可选择由 conda 管理的 cudatoolkit,通常无需单独安装系统级 CUDA 与 cuDNN;若使用 pip 且希望使用系统 CUDA,则需按上文安装并配置环境变量。
三 安装支持 GPU 的 PyTorch
- 使用 pip(按 CUDA 版本选择命令)
- CUDA 11.7 示例:pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
- CUDA 11.3 示例:pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
- 其他版本:将 cu 替换为你的 CUDA 主版本(如 cu118 等),命令模板:
- pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu
- 使用 conda(由 conda 管理 cudatoolkit)
- CUDA 11.7 示例:conda install pytorch torchvision torchaudio cudatoolkit=11.7 -c pytorch
- 提示
- 访问 pytorch.org/get-started/locally 可获取与你的系统与 CUDA 版本匹配的最新安装命令。
四 验证与常见问题
- 验证安装
- 执行:python - <<‘PY’
import torch
print(“torch:”, torch.version)
print(“cuda available:”, torch.cuda.is_available())
print(“device count:”, torch.cuda.device_count())
if torch.cuda.is_available():
print(“current device:”, torch.cuda.current_device())
print(“device name:”, torch.cuda.get_device_name(torch.cuda.current_device()))
PY
- 预期输出:cuda available: True、显示 GPU 数量与名称。
- 常见问题排查
- torch.cuda.is_available() 为 False
- 检查 nvidia-smi 是否正常;若异常,重装/修复 NVIDIA 驱动(如 ubuntu-drivers autoinstall 后重启)。
- 核对 CUDA 与 PyTorch 版本匹配;若使用 pip 且依赖系统 CUDA,确认 PATH/LD_LIBRARY_PATH 包含 /usr/local/cuda-/bin 与 /usr/local/cuda-/lib64。
- 若使用 conda,优先用 cudatoolkit 方式避免系统库冲突。
- 版本对应关系
- 注意 torch 与 torchvision 的版本对应关系;如采用本地 whl 安装,请严格匹配版本与 CUDA 标签(如 cu113 对应 CUDA 11.3)。