Ubuntu上PyTorch与CUDA的集成指南

Ubuntu上PyTorch与CUDA集成指南

一 环境准备与兼容性检查

• 确认硬件与系统：需要NVIDIA GPU且支持CUDA。在终端执行nvidia-smi查看驱动与GPU信息；执行nvcc --version查看CUDA编译器版本。若nvidia-smi显示的CUDA Runtime与nvcc显示的CUDA Toolkit不一致，这是常见现象，通常以PyTorch预编译包的CUDA版本为准。建议优先使用Ubuntu 20.04/22.04 LTS以获得更好的兼容性与驱动支持。为降低冲突，建议使用conda或venv创建隔离环境。

二 安装驱动与CUDA Toolkit

• 安装驱动（二选一）：
  • 自动安装推荐驱动：sudo ubuntu-drivers autoinstall
  • 指定版本（示例）：sudo apt install nvidia-driver-555
• 安装CUDA Toolkit（两种常用方式，择一）：
  • 使用官方网络仓库安装（推荐，便于升级）：
    • Ubuntu 20.04示例：

      wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb
      sudo dpkg -i cuda-keyring_1.0-1_all.deb
      sudo apt update
      sudo apt install cuda-toolkit-12-1
      

  • 使用官方runfile安装（适合需要自定义组件的场景）：

    wget https://developer.download.nvidia.com/compute/cuda/<
        version>
        /cuda_<
        version>
        _linux.run
    sudo sh cuda_<
        version>
        _linux.run
    

• 环境变量（安装完成后执行一次即可，路径按实际调整）：

  echo 'export PATH=/usr/local/cuda/bin:$PATH' >
      >
       ~/.bashrc
  echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >
      >
       ~/.bashrc
  source ~/.bashrc
  

• 验证：nvidia-smi应显示GPU与驱动信息，nvcc --version应显示Toolkit版本。

三 安装cuDNN

• 下载与CUDA版本匹配的cuDNN（需NVIDIA开发者账号），解压并按如下方式部署（示例路径为/usr/local/cuda）：

  tar -xzvf cudnn-<
      version>
      -linux-x64-v<
      version>
      .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*
  

• 注意：如果使用APT方式安装CUDA，也可通过对应仓库安装匹配的libcudnn包；务必保持与CUDA版本一致。

四 安装PyTorch GPU版本

• 选择安装方式（pip或conda），确保与系统的CUDA大版本匹配（如本地为CUDA 12.x，优先选择cu12x系列；如为11.x，选择cu11x系列）。
• pip安装示例（按你的CUDA版本替换cu118/cu121/cu124等）：

  pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
  

• conda安装示例（conda会自动携带对应的cudatoolkit，无需单独安装系统级CUDA）：

  conda create -n pytorch_env python=3.10 -y
  conda activate pytorch_env
  conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
  

• 说明：PyTorch与CUDA不需要“版本号完全一致”，但需大版本匹配；若本地为CUDA 12.4，通常选择cu121/cu124的PyTorch预编译包均可稳定运行。

五 验证与常见问题处理

• 验证脚本：

  python - <
      <
      'PY'
  import torch
  print("PyTorch:", torch.__version__)
  print("CUDA可用:", torch.cuda.is_available())
  print("CUDA版本:", torch.version.cuda)
  print("设备数:", torch.cuda.device_count())
  if torch.cuda.is_available():
      print("设备名:", torch.cuda.get_device_name(0))
  PY
  

  预期输出包含：True、正确的CUDA版本号与GPU型号。
• 常见问题与排查要点：
  • torch.cuda.is_available()返回False：
    • 检查驱动：nvidia-smi是否正常；必要时重装驱动或重启。
    • 检查环境变量：PATH与LD_LIBRARY_PATH是否包含CUDA的bin与lib64。
    • 检查cuDNN：头文件与库文件是否正确部署且权限正确。
    • 检查PyTorch与CUDA匹配：优先使用与本地CUDA 大版本一致的PyTorch预编译包。
  • 编译扩展（如flash-attn）失败：多由PyTorch与CUDA版本不匹配引起，先按上文原则统一版本后重装扩展。

声明：本文内容由网友自发贡献，本站不承担相应法律责任。对本内容有异议或投诉，请联系2913721942#qq.com核实处理，我们将尽快回复您，谢谢合作！