PyTorch Ubuntu环境怎么配置

PyTorch在Ubuntu环境中的配置步骤

1. 更新系统包

首先确保Ubuntu系统包是最新的，避免依赖冲突：

sudo apt update && sudo apt upgrade -y

2. 安装Python与pip

PyTorch需要Python 3.6及以上版本，Ubuntu默认仓库中的Python 3已满足要求：

sudo apt install -y python3 python3-pip

3. 创建虚拟环境（推荐）

为隔离项目依赖，避免包冲突，建议创建虚拟环境：

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

4. 安装PyTorch

PyTorch提供CPU版本（无GPU加速）和GPU版本（需CUDA支持），根据需求选择：

（1）CPU版本（无GPU）

直接通过pip安装预编译的CPU版本：

pip install torch torchvision torchaudio

（2）GPU版本（需NVIDIA GPU）

需先确保系统安装了兼容的CUDA Toolkit和cuDNN（详见下文“5. 安装CUDA与cuDNN”），然后通过pip安装对应版本的PyTorch：

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118  # 以CUDA 11.8为例

注：替换cu118为你的CUDA版本（如cu117、cu121），可通过nvcc --version查看已安装的CUDA版本。

5. 安装CUDA与cuDNN（GPU版本必需）

若需使用GPU加速，需完成以下步骤：

（1）安装NVIDIA驱动

sudo apt install -y nvidia-driver-535  # 替换为你的显卡支持的驱动版本（可通过`ubuntu-drivers devices`查看推荐）
sudo reboot  # 重启系统使驱动生效

（2）安装CUDA Toolkit

以CUDA 11.8为例，通过NVIDIA官方仓库安装：

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/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-11.8.0-520.61.05-1/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt update
sudo apt install -y cuda-11-8  # 安装CUDA 11.8

（3）安装cuDNN

cuDNN是NVIDIA的深度学习加速库，需注册NVIDIA开发者账号下载：

wget https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.9.7.29/Production/11.8_20230725/cudnn-linux-x86_64-8.9.7.29_cuda11-archive.tar.xz
tar -xf cudnn-linux-x86_64-8.9.7.29_cuda11-archive.tar.xz
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda/include
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda/lib64/
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

（4）配置环境变量

编辑~/.bashrc文件，添加CUDA路径：

echo 'export PATH=/usr/local/cuda-11.8/bin${PATH:+:${PATH}}' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc
source ~/.bashrc  # 使环境变量生效

6. 验证PyTorch安装

在Python交互环境中运行以下代码，检查PyTorch版本及CUDA是否可用：

import torch
print("PyTorch版本:", torch.__version__)
print("CUDA是否可用:", torch.cuda.is_available())  # 若为True则表示GPU版本安装成功

若输出类似以下内容，则说明安装成功：

PyTorch版本: 2.1.0
CUDA是否可用: True

7. 安装其他依赖库（可选）

根据项目需求，可安装常用的Python库：

pip install numpy pandas matplotlib scikit-learn

8. 配置IDE（可选）

若使用PyCharm或VSCode，需将虚拟环境的Python解释器添加到IDE中：

• PyCharm：File → Settings → Project → Python Interpreter → 添加解释器，选择虚拟环境中的pytorch_env/bin/python。
• VSCode：File → Preferences → Settings → Python: Python Path，设置为虚拟环境的解释器路径。

通过以上步骤，你即可在Ubuntu环境中成功配置PyTorch，满足CPU或GPU加速的深度学习需求。