Ubuntu如何配置PyTorch的网络环境

Ubuntu配置PyTorch网络环境的详细步骤

1. 安装系统依赖

首先更新系统包并安装Python基础环境，确保后续安装顺利：

sudo apt update && sudo apt upgrade -y
sudo apt install python3 python3-pip python3-venv -y

2. 安装NVIDIA显卡驱动

PyTorch GPU加速依赖NVIDIA显卡驱动，需先确认显卡型号并安装对应驱动：

查看显卡型号：lspci | grep -i nvidia
安装推荐驱动：运行ubuntu-drivers devices获取推荐驱动版本，再用以下命令安装（以nvidia-driver-535为例）：
```
sudo apt install nvidia-driver-535 -y
sudo reboot  # 重启系统使驱动生效
```
验证驱动：nvidia-smi（显示显卡信息及驱动版本即为成功）。

3. 安装CUDA Toolkit

CUDA是NVIDIA的并行计算平台，PyTorch GPU版本需匹配对应CUDA版本：

下载CUDA：前往NVIDIA CUDA Toolkit官网，选择Ubuntu版本（如22.04）及CUDA版本（如11.8），下载.run或.deb安装包。

安装CUDA（以.deb包为例）：

wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt update
sudo apt install cuda-toolkit-11-8 -y

配置环境变量：编辑~/.bashrc，添加以下内容：

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

运行source ~/.bashrc使变量生效。

验证CUDA：nvcc --version（显示CUDA编译器版本即为成功）。

4. 安装cuDNN库

cuDNN是NVIDIA的深度学习加速库，需与CUDA版本匹配：

下载cuDNN：前往NVIDIA cuDNN官网，登录后下载对应CUDA版本（如CUDA 11.8）的cuDNN包（如cudnn-linux-x86_64-8.6.0.163_cuda11-archive.tar.xz）。

安装cuDNN：

tar -xvf cudnn-linux-x86_64-8.6.0.163_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*

验证cuDNN：ls /usr/local/cuda/lib64 | grep libcudnn（显示libcudnn文件即为成功）。

5. 创建虚拟环境（可选但推荐）

虚拟环境可隔离PyTorch与其他项目的依赖冲突：

python3 -m venv pytorch_env
source pytorch_env/bin/activate  # 激活环境（后续操作均在激活状态下进行）

6. 安装PyTorch

根据是否需要GPU加速，选择以下安装方式：

① GPU版本（推荐）

通过PyTorch官网获取对应CUDA版本的安装命令（如CUDA 11.8）：

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

② CPU版本

若无需GPU加速，安装CPU版本：

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

7. 验证安装

运行以下Python代码，检查PyTorch版本及CUDA是否可用：

import torch
print("PyTorch版本:", torch.__version__)
print("CUDA是否可用:", torch.cuda.is_available())  # 返回True表示GPU可用
print("GPU设备名称:", torch.cuda.get_device_name(0))  # 显示GPU型号

常见问题解决

nvcc: command not found：CUDA未正确安装或环境变量未配置，检查~/.bashrc中的路径是否包含CUDA的bin目录。
torch.cuda.is_available()返回False：可能是CUDA/cuDNN版本与PyTorch不匹配，或驱动未正确安装，参考PyTorch官网版本兼容表调整。

网络安装慢：使用国内镜像源（如清华源），修改pip命令为：

pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple --extra-index-url https://download.pytorch.org/whl/cu118

通过以上步骤，即可在Ubuntu系统上完成PyTorch的网络环境配置，支持CPU或GPU加速。如需进一步优化（如Jupyter Notebook集成），可参考PyTorch官方文档。