Ubuntu上如何搭建PyTorch开发环境

Ubuntu 搭建 PyTorch 开发环境

一 准备与系统检查

• 更新系统并安装基础工具：
  • sudo apt update && sudo apt upgrade -y
  • sudo apt install -y build-essential cmake git wget unzip yasm pkg-config libopenblas-dev liblapack-dev libjpeg-dev libpng-dev python3 python3-pip
• 确认 GPU 与驱动：
  • lspci | grep -i nvidia 查看是否有 NVIDIA 显卡
  • nvidia-smi 查看驱动与 CUDA 运行时版本（右上角显示 Supported/Runtime CUDA）
• 选择安装方式：优先使用 Miniconda/Anaconda 创建隔离环境；GPU 训练需安装与 PyTorch 版本匹配的 CUDA/cuDNN，或直接使用 PyTorch 官方预编译包（自带 CUDA 运行时）。

二 安装 Miniconda 与创建虚拟环境

• 安装 Miniconda（示例为 x86_64）：
  • wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
  • bash Miniconda3-latest-Linux-x86_64.sh（按提示完成安装）
• 初始化与常用配置：
  • conda init bash && source ~/.bashrc
  • conda config --set always_yes yes --set changeps1 no
  • conda config --add channels defaults --add channels conda-forge --add channels r
• 创建并激活环境（示例：Python 3.10）：
  • conda create -n pytorch-env python=3.10
  • conda activate pytorch-env
• 说明：也可使用 Anaconda；在 conda 环境中安装 PyTorch 更便于依赖隔离与多版本管理。

三 安装 PyTorch 与 CUDA 方案

• GPU 方案 A（推荐，使用官方预编译包，含 CUDA 运行时）：
  • 在激活的 conda 环境中执行（示例为 CUDA 11.8；请按官网生成命令替换版本）：
    • pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
• GPU 方案 B（系统级 CUDA + cuDNN，适合需要系统级库的场景）：
  • 安装与驱动匹配的 CUDA（示例 11.3）：
    • 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.3.1/local_installers/cuda-repo-ubuntu2204-11-3-local-11.3.1-450.51.06-1_amd64.deb
    • sudo dpkg -i cuda-repo-ubuntu2204-11-3-local-11.3.1-450.51.06-1_amd64.deb
    • sudo apt-key add /var/cuda-repo-ubuntu2204-11-3-local-11.3.1-450.51.06-1/7fa2af80.pub
    • sudo apt update && sudo apt install -y cuda
  • 安装与 CUDA 匹配的 cuDNN（示例 8.0）：
    • 从 NVIDIA 下载 cuDNN 对应版本（需登录），解压后复制头文件与库：
      • tar -xzvf cudnn-archive-version-linux-x64-v8.0.3.33.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
  • 安装 PyTorch（与系统 CUDA 版本匹配，示例 cu113）：
    • pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
• CPU 方案（无 NVIDIA GPU 或仅 CPU 训练）：
  • pip install torch torchvision torchaudio
• 提示：GPU 方案 A 更省事；方案 B 更灵活但依赖管理更复杂。安装命令中的 CUDA 版本需与 PyTorch 版本匹配，优先参考官网生成命令。

四 验证安装与常见问题

• 验证安装：
  • 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
• 常见问题与处理：
  • 版本不匹配：nvidia-smi 显示的 CUDA 为驱动支持上限；PyTorch 二进制包自带 CUDA Runtime。若报错，请更换与 PyTorch 一致的 CUDA 标签（如 cu118）。
  • 网络慢/超时：pip 换国内镜像（如清华源）或使用 conda 官方通道；必要时重试或使用离线安装包。
  • 权限问题：避免使用 sudo pip；优先在用户环境或 conda 环境中安装。
  • 驱动问题：使用 sudo ubuntu-drivers autoinstall 安装推荐驱动，或到 NVIDIA 官网下载对应版本。