在 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）。