Ubuntu下如何搭建PyTorch环境

Ubuntu下搭建PyTorch环境的详细步骤

1. 更新系统包

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

sudo apt update &
    &
 sudo apt upgrade -y

2. 安装必要依赖

PyTorch及后续工具（如CUDA、conda）的安装需要基础开发工具和库：

sudo apt install -y build-essential cmake git wget unzip yasm pkg-config \
    libopenblas-dev liblapack-dev libjpeg-dev libpng-dev

3. 安装Python与pip

PyTorch依赖Python 3.6及以上版本，Ubuntu默认仓库的Python 3通常满足要求：

sudo apt install -y python3 python3-pip
# 验证安装
python3 --version  # 应输出3.6及以上版本
pip3 --version     # 应输出pip版本

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

虚拟环境可隔离项目依赖，避免与系统或其他项目冲突：

# 安装venv模块（若未安装）
sudo apt install -y python3-venv
# 创建虚拟环境（名称自定义，如pytorch_env）
python3 -m venv pytorch_env
# 激活虚拟环境
source pytorch_env/bin/activate
# 激活后，终端提示符会显示环境名称（如(pytorch_env) user@ubuntu:~$）

5. 安装CUDA与cuDNN（可选，GPU加速必需）

若需使用GPU加速PyTorch，需安装NVIDIA CUDA Toolkit和cuDNN库。以下以CUDA 11.8为例（需根据显卡型号选择兼容版本，如RTX 30系列支持CUDA 11.8+）：

5.1 安装NVIDIA显卡驱动

# 自动检测并安装适合显卡的最新驱动
sudo ubuntu-drivers autoinstall
# 重启系统使驱动生效
sudo reboot
# 验证驱动安装
nvidia-smi  # 应显示显卡型号及驱动版本

5.2 安装CUDA Toolkit

# 下载CUDA 11.8的.deb安装包（替换为最新版本链接）
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 apt-key add /var/cuda-repo-ubuntu2204-11-8-local/7fa2af80.pub
sudo apt update
sudo apt install -y cuda
# 添加CUDA到环境变量（编辑~/.bashrc）
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
# 验证CUDA安装
nvcc --version  # 应显示CUDA 11.8版本

5.3 安装cuDNN

# 注册NVIDIA开发者账号（免费），下载对应CUDA版本的cuDNN（如CUDA 11.8对应cuDNN 8.9.7）
# 解压并复制文件到CUDA目录
tar -xzvf cudnn-linux-x64-v8.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*
# 验证cuDNN安装
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2  # 应显示cuDNN版本

6. 安装PyTorch

通过pip或conda安装PyTorch，根据是否使用GPU选择对应命令：

6.1 使用pip安装

• CPU版本（无GPU加速）：

  pip install torch torchvision torchaudio
  

• GPU版本（需匹配CUDA版本，如CUDA 11.8）：

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

6.2 使用conda安装（推荐，环境管理更方便）

• CPU版本：

  conda install pytorch torchvision torchaudio cpuonly -c pytorch
  

• GPU版本（如CUDA 11.8）：

  conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
  

7. 验证安装

安装完成后，通过Python代码验证PyTorch是否正常工作及是否支持GPU：

import torch
print("PyTorch版本:", torch.__version__)  # 应输出安装的PyTorch版本
print("CUDA是否可用:", torch.cuda.is_available())  # 若为True，说明GPU加速可用
print("GPU设备名称:", torch.cuda.get_device_name(0))  # 显示GPU型号（如NVIDIA RTX 3090）

若torch.cuda.is_available()返回True，则表示GPU加速配置成功。

注意事项

• 版本兼容性：PyTorch版本需与CUDA版本匹配（如PyTorch 2.0+支持CUDA 11.8+），可通过PyTorch官方兼容性表确认。
• 虚拟环境：始终在虚拟环境中安装PyTorch，避免依赖冲突。
• 驱动更新：定期更新NVIDIA驱动以支持最新CUDA版本。

通过以上步骤，即可在Ubuntu系统上成功搭建PyTorch环境，满足深度学习项目的需求。

声明：本文内容由网友自发贡献，本站不承担相应法律责任。对本内容有异议或投诉，请联系2913721942#qq.com核实处理，我们将尽快回复您，谢谢合作！