Ubuntu中PyTorch的安装教程

Ubuntu系统PyTorch安装详细教程

一、前置准备

在安装PyTorch前，需完成以下基础环境配置：

更新系统包
运行以下命令确保系统软件包为最新版本：
```
sudo apt update && sudo apt upgrade -y
```
安装Python及pip
Ubuntu 20.04及以上版本默认包含Python 3.8+，需确认安装python3、pip及虚拟环境工具：
```
sudo apt install python3 python3-pip python3-venv -y
```

创建虚拟环境（推荐）
为避免项目依赖冲突，建议使用虚拟环境隔离PyTorch及其依赖：

python3 -m venv pytorch_env  # 创建名为pytorch_env的虚拟环境
source pytorch_env/bin/activate  # 激活虚拟环境（激活后终端提示符会显示环境名）

激活后，后续所有安装操作均在虚拟环境中进行。

二、安装PyTorch（分GPU/CPU版本）

1. GPU版本（需NVIDIA显卡支持）

若需使用GPU加速，需先安装兼容的CUDA Toolkit和cuDNN库（具体版本需匹配PyTorch版本，参考PyTorch官网要求）。以下以Ubuntu 22.04 + CUDA 12.6为例：

步骤1：安装CUDA Toolkit
访问NVIDIA CUDA Toolkit官网下载对应Ubuntu版本的安装包（如CUDA 12.6），或通过以下命令安装：

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 https://developer.download.nvidia.com/compute/cuda/12.6.0/local_installers/cuda-repo-ubuntu2204-12-6-local_12.6.0-535.104.53-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-12-6-local_12.6.0-535.104.53-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-12-6-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt update
sudo apt install cuda -y

安装完成后，需将CUDA路径添加至环境变量（编辑~/.bashrc文件）：

echo 'export PATH=/usr/local/cuda-12.6/bin${PATH:+:${PATH}}' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.6/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc
source ~/.bashrc

步骤2：安装cuDNN库
访问NVIDIA cuDNN官网（需注册开发者账号）下载对应CUDA版本的cuDNN（如CUDA 12.6对应cuDNN 8.9.5），解压后复制文件至CUDA目录：

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

步骤3：安装PyTorch GPU版本
通过pip安装与CUDA版本匹配的PyTorch预编译包（以CUDA 12.6为例）：

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

若使用conda安装，命令如下：

conda install pytorch torchvision torchaudio cudatoolkit=12.6 -c pytorch -c conda-forge

2. CPU版本（无GPU加速）

若无需GPU加速，可直接安装CPU版本的PyTorch（安装速度更快，无需CUDA依赖）：

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

或使用conda安装：

conda install pytorch torchvision torchaudio cpuonly -c pytorch

三、验证安装

安装完成后，通过以下Python代码验证PyTorch是否安装成功及GPU是否可用：

import torch
print("PyTorch版本:", torch.__version__)  # 输出PyTorch版本号
print("CUDA可用性:", torch.cuda.is_available())  # 若为True则表示GPU可用
if torch.cuda.is_available():
    print("CUDA设备名称:", torch.cuda.get_device_name(0))  # 输出GPU设备名称

若输出显示PyTorch版本号且torch.cuda.is_available()返回True（GPU版本），则说明安装成功。

四、常见问题解决

CUDA驱动不兼容：若安装GPU版本后出现CUDA error: no kernel image is available for execution on the device错误，需升级NVIDIA驱动至与CUDA版本匹配（参考NVIDIA官方驱动下载页面）。
虚拟环境未激活：若安装后无法导入PyTorch，需确认是否激活了虚拟环境（source pytorch_env/bin/activate）。
安装速度慢：可使用国内镜像源加速pip安装（如pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch），但需注意版本兼容性。