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 # 激活虚拟环境(激活后终端提示符会显示环境名)
激活后,后续所有安装操作均在虚拟环境中进行。
若需使用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
若无需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 error: no kernel image is available for execution on the device错误,需升级NVIDIA驱动至与CUDA版本匹配(参考NVIDIA官方驱动下载页面)。source pytorch_env/bin/activate)。pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch),但需注意版本兼容性。