Ubuntu下如何安装PyTorch的GPU版本

Ubuntu下安装PyTorch GPU版本的完整步骤

1. 前置准备：更新系统及安装基础依赖

在安装前，确保系统包列表是最新的，并安装构建工具与Python环境：

sudo apt update &
    &
 sudo apt upgrade -y
sudo apt install -y python3 python3-pip python3-venv build-essential cmake git wget unzip yasm pkg-config libopenblas-dev liblapack-dev libjpeg-dev libpng-dev

这一步为后续安装驱动、CUDA及PyTorch提供基础支持。

2. 安装NVIDIA显卡驱动

PyTorch GPU版本依赖NVIDIA驱动，需先确认驱动是否安装：

nvidia-smi

若命令输出GPU信息（如型号、驱动版本），则驱动已安装；若未安装，可通过以下方式安装：

• 自动安装推荐驱动（推荐）：

  sudo ubuntu-drivers autoinstall
  

  该命令会根据显卡型号自动安装兼容的最新驱动，安装完成后重启系统：

  sudo reboot
  

• 手动安装指定驱动（可选）：
  若需安装特定版本驱动（如470），可使用：

  sudo apt install -y nvidia-driver-470
  sudo reboot
  

  重启后再次运行nvidia-smi确认驱动生效。

3. 安装CUDA Toolkit

CUDA是PyTorch调用GPU的核心组件，需选择与PyTorch版本兼容的CUDA版本（如PyTorch 2.6.0+对应CUDA 12.6）。以下以CUDA 12.6为例：

• 下载CUDA Toolkit：
  访问NVIDIA CUDA Toolkit下载页面，选择Ubuntu版本（如22.04），下载对应的.deb文件（如cuda-repo-ubuntu2204-12-6-local_12.6.0-530.30.02-1_amd64.deb）。
• 安装CUDA：

  sudo dpkg -i cuda-repo-ubuntu2204-12-6-local_12.6.0-530.30.02-1_amd64.deb
  sudo cp /var/cuda-repo-ubuntu2204-12-6-local/cuda-*-keyring.gpg /usr/share/keyrings/
  sudo apt update
  sudo apt install -y cuda
  

• 设置环境变量：
  编辑~/.bashrc文件，添加以下内容：

  export PATH=/usr/local/cuda/bin${
  PATH:+:${
  PATH}
  }
  
  export LD_LIBRARY_PATH=/usr/local/cuda/lib64${
  LD_LIBRARY_PATH:+:${
  LD_LIBRARY_PATH}
  }
      
  

  使环境变量生效：

  source ~/.bashrc
  

• 验证CUDA安装：
  运行nvcc --version，若输出CUDA版本信息（如release 12.6），则安装成功。

4. 安装cuDNN库

cuDNN是深度学习专用加速库，需与CUDA版本匹配（如CUDA 12.6对应cuDNN 8.9+）。安装步骤：

• 下载cuDNN：
  访问NVIDIA cuDNN下载页面，注册并登录账号，选择与CUDA版本匹配的cuDNN（如cuDNN Library for Linux），下载.tgz文件。
• 安装cuDNN：

  tar -xzvf cudnn-12.6-linux-x64-v8.9.7.29.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*
  

  这一步将cuDNN库文件复制到CUDA目录，供PyTorch调用。

5. 安装PyTorch GPU版本

通过pip安装PyTorch GPU版本（以PyTorch 2.6.0+cu126为例）：

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

• 版本匹配说明：
  PyTorch版本需与CUDA版本严格对应（如PyTorch 2.1.0对应CUDA 12.1，PyTorch 2.6.0对应CUDA 12.6），具体版本可参考PyTorch官方安装指南。

6. 验证安装

运行以下Python代码，检查PyTorch是否成功识别GPU：

import torch
print("PyTorch版本:", torch.__version__)
print("CUDA可用性:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("CUDA版本:", torch.version.cuda())
    print("GPU设备名称:", torch.cuda.get_device_name(0))

若输出中torch.cuda.is_available()为True，且显示GPU设备名称（如NVIDIA GeForce RTX 3060），则说明安装成功。

注意事项

• 驱动与CUDA版本兼容性：确保驱动版本支持已安装的CUDA版本（如CUDA 12.6需驱动版本≥515.65.01），可通过NVIDIA官方文档查询兼容列表。
• 虚拟环境推荐：建议使用venv或conda创建虚拟环境，避免依赖冲突（如python3 -m venv pytorch_gpu_env & & source pytorch_gpu_env/bin/activate）。
• 版本更新：定期检查PyTorch、CUDA及驱动的更新，以获取最新性能优化和安全补丁。

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