Ubuntu PyTorch环境配置技巧

Ubuntu下PyTorch环境配置技巧

1. 系统准备：更新与依赖安装

在配置前，先更新系统软件包以确保兼容性：

sudo apt update && sudo apt upgrade -y

安装Python基础依赖（Python3、pip、构建工具）：

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

这些依赖是PyTorch及后续扩展库（如OpenCV）正常运行的基础。

2. 虚拟环境：隔离项目依赖

使用venv或conda创建隔离环境，避免全局包冲突：

venv（Python内置）：

python3 -m venv pytorch_env  # 创建环境
source pytorch_env/bin/activate  # 激活环境（提示符变为“(pytorch_env)”）

conda（推荐，更适合复杂项目）：

wget https://repo.anaconda.com/archive/Anaconda3-2025.02-Linux-x86_64.sh  # 下载Anaconda
bash Anaconda3-2025.02-Linux-x86_64.sh  # 按提示安装
conda create -n pytorch_env python=3.9 -y  # 创建环境（指定Python版本）
conda activate pytorch_env  # 激活环境

虚拟环境能有效隔离不同项目的依赖版本，防止冲突。

3. PyTorch安装：选择CPU/GPU版本

3.1 CPU版本（无GPU加速）

通过pip安装CPU版本（适合入门或无GPU设备）：

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

3.2 GPU版本（需NVIDIA显卡）

需先安装CUDA Toolkit和cuDNN，再安装对应版本的PyTorch：

步骤1：安装CUDA Toolkit
从NVIDIA官网下载与显卡兼容的CUDA版本（如CUDA 11.7），通过.deb文件安装：

sudo dpkg -i cuda-repo-ubuntu2204-11-7-local_11.7.99-515.65.01-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt update
sudo apt install -y cuda

步骤2：安装cuDNN
注册NVIDIA开发者账号，下载与CUDA版本匹配的cuDNN（如cuDNN 8.9.5 for CUDA 11.7），解压后复制到CUDA目录：

tar -xzvf cudnn-linux-x86_64-8.9.5.29_cuda11-archive.tar.xz
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
通过pip安装对应CUDA版本的PyTorch（如CUDA 11.7）：
```
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
```
或使用conda（更便捷，自动处理依赖）：
```
conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
```
安装后通过torch.cuda.is_available()验证GPU是否可用。

4. 版本兼容性：避免踩坑

PyTorch与CUDA版本：参考PyTorch官网兼容性表（如PyTorch 2.0支持CUDA 11.7/11.8，PyTorch 1.13支持CUDA 11.6/11.7）。可通过nvcc --version查看已安装CUDA版本，选择匹配的PyTorch版本。
Ubuntu与CUDA版本：Ubuntu 22.04 LTS支持CUDA 12.x及以上，Ubuntu 20.04 LTS支持CUDA 11.x及以下，选择与系统匹配的CUDA版本以避免驱动冲突。

5. 验证安装：确认环境正常

运行以下Python代码验证PyTorch安装及GPU支持：

import torch
print("PyTorch版本:", torch.__version__)  # 查看PyTorch版本
print("CUDA可用性:", torch.cuda.is_available())  # 检查GPU是否可用
if torch.cuda.is_available():
    print("GPU设备名称:", torch.cuda.get_device_name(0))  # 查看GPU型号

若输出显示CUDA可用且GPU型号正确，则配置成功。

6. 常见问题解决技巧

CUDA驱动问题：若nvidia-smi显示驱动版本过低，可通过sudo ubuntu-drivers autoinstall自动安装最新驱动，或手动下载对应版本的驱动（如nvidia-driver-555）安装。
环境变量配置：若CUDA未正确识别，可编辑~/.bashrc添加以下内容（替换为实际CUDA路径）：
```
export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
```
运行source ~/.bashrc使配置生效。
依赖冲突：若安装过程中出现依赖错误，可使用conda clean --all清理conda缓存，或pip cache purge清理pip缓存后重新安装。