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及驱动的更新，以获取最新性能优化和安全补丁。