PyTorch在Ubuntu上有GPU加速吗

支持情况与前提条件

• 在Ubuntu上，PyTorch可以使用NVIDIA GPU进行加速，需要同时满足：
  • 硬件为NVIDIA 显卡且支持CUDA。
  • 正确安装NVIDIA 显卡驱动。
  • 安装与PyTorch版本匹配的CUDA/cuDNN（或使用PyTorch官方预编译的CUDA包）。
  • 安装带有CUDA支持的PyTorch发行版（pip/conda）。

快速验证是否可用

• 查看驱动与GPU：
  • 运行：nvidia-smi，应显示GPU型号、驱动版本、CUDA版本等信息。
• 在Python中检查CUDA可用性并查看设备信息：
  • 代码示例：
    • import torch; print(torch.cuda.is_available())
    • print(torch.cuda.device_count())
    • print(torch.cuda.get_device_name(0))
  • 若torch.cuda.is_available()返回True，说明PyTorch已成功启用GPU加速。

安装要点

• 安装或确认NVIDIA驱动：
  • 可自动安装推荐驱动：sudo ubuntu-drivers autoinstall && sudo reboot
  • 或指定版本：sudo apt install nvidia-driver-<版本号>。
• 安装CUDA/cuDNN（两种常见做法）：
  • 系统级安装：安装与驱动匹配的CUDA Toolkit与cuDNN库（注意版本兼容）。
  • 使用PyTorch官方预编译包：通过pip/conda安装带cu前缀的PyTorch，无需单独安装系统CUDA（例如：pip install torch ... --extra-index-url https://download.pytorch.org/whl/cu118 或 conda install pytorch cudatoolkit=11.8 -c pytorch）。
• 安装PyTorch GPU版（示例命令，按你的CUDA版本选择其一）：
  • pip：pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
  • conda：conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch
  • 安装后用前述Python代码验证torch.cuda.is_available()。

常见故障排查

• nvidia-smi不可用或报错：重装/更新驱动，使用ubuntu-drivers autoinstall后重启。
• torch.cuda.is_available()返回False：
  • 检查驱动与CUDA/cuDNN版本是否匹配。
  • 确认安装的是GPU版PyTorch而非CPU版。
  • 若采用系统CUDA，核对环境变量（如PATH、LD_LIBRARY_PATH）是否包含CUDA路径。
  • 确保Python环境中没有多个冲突的PyTorch安装。