温馨提示×

Ubuntu PyTorch与CUDA怎么安装

小樊
37
2025-12-28 12:49:05
栏目: 智能运维

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

0