Ubuntu 20.04 完美安装 PyTorch 2.7.1 指南：避坑与优化

内容摘要：Ubuntu 20.04 完美安装 PyTorch 2.7.1 指南：避坑与优化,

在 Ubuntu 20.04 系统上安装 PyTorch 2.1.0 并非总是一帆风顺。很多开发者，特别是依赖深度学习框架进行模型训练或推理的同学，经常会遇到 CUDA 版本不兼容、依赖包缺失、甚至由于网络问题导致下载失败等问题。本文将详细介绍如何在 Ubuntu 20.04 上成功安装 PyTorch 2.1.0，并分享一些常见的坑以及相应的解决方案。

1. 环境准备：CUDA、cuDNN 与驱动版本

在安装 PyTorch 之前，我们需要确保 CUDA、cuDNN 以及 NVIDIA 驱动程序版本兼容。PyTorch 官方文档会明确指出兼容的 CUDA 版本。不匹配的版本会导致运行时错误，比如常见的 CUDA driver version is insufficient for CUDA runtime version 错误。

查看 NVIDIA 驱动版本：

nvidia-smi

输出结果会显示驱动版本和 CUDA 版本。如果没有安装 NVIDIA 驱动，可以使用以下命令安装：

sudo apt update
sudo apt install nvidia-driver-XXX  # XXX 替换为推荐的版本，例如 nvidia-driver-535

安装 CUDA：

下载对应版本的 CUDA Toolkit。建议从 NVIDIA 官网下载 .deb 包，方便使用 apt 管理。

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
sudo apt update
sudo apt install cuda

安装完成后，配置环境变量：

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

将上述环境变量添加到 ~/.bashrc 或者 ~/.zshrc 文件中，并执行 source ~/.bashrc 或 source ~/.zshrc 使其生效。

安装 cuDNN：

cuDNN 通常需要从 NVIDIA 开发者网站下载，并且需要注册账号。下载对应 CUDA 版本的 cuDNN 包，解压后将文件复制到 CUDA 安装目录：

tar -xzv cudnn-linux-x86_64-8.x.x.x_cudaX.Y-archive.tar.gz
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*

2. 安装 PyTorch 2.1.0

建议使用 Anaconda 或 Miniconda 创建一个独立的虚拟环境，避免与其他 Python 包产生冲突。

conda create -n pytorch python=3.8  # 建议Python版本为3.8或3.9
conda activate pytorch

使用 pip 安装 PyTorch，可以指定 CUDA 版本。PyTorch 官网会提供安装命令。

pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118 # cu118 表示 CUDA 11.8

如果国内网络环境较差，可以考虑使用国内镜像源，例如清华大学、上海交通大学等。例如：

pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 验证安装

在 Python 交互环境中，验证 PyTorch 是否成功安装并且能够使用 CUDA：

import torch
print(torch.__version__)
print(torch.cuda.is_available())

如果 torch.cuda.is_available() 返回 True，则表示 PyTorch 已经成功安装并且可以利用 CUDA 进行 GPU 加速。

4. 常见问题与解决方案

• CUDA 版本不匹配： 确保 PyTorch 版本与 CUDA 版本兼容。参考 PyTorch 官方文档。
• 依赖包缺失： 安装 PyTorch 过程中可能会提示缺少依赖包，例如 ninja。使用 pip install ninja 安装即可。
• 网络问题： 如果下载速度慢或下载失败，尝试更换国内镜像源，或者使用代理。
• RuntimeError: cuDNN error： 可能是 cuDNN 安装不正确，或者版本不兼容。重新安装 cuDNN，并确保文件复制到正确的目录。

5. 性能优化建议

• 使用 GPU 进行模型训练： 确保模型和数据都加载到 GPU 上。
• 使用 torch.backends.cudnn.benchmark = True： 允许 cuDNN 寻找最佳的卷积算法，提高训练速度。
• 合理设置 batch size： 根据 GPU 显存大小调整 batch size，充分利用 GPU 资源。

Ubuntu 20.04下的Pytorch2.1.0安装总结

希望通过本文，你能够顺利在 Ubuntu 20.04 系统上安装 PyTorch 2.1.0。记住，遇到问题时，仔细阅读错误信息，查阅官方文档和社区论坛，通常都能找到解决方案。善用搜索引擎，例如 Google 和 Stack Overflow，也是解决问题的关键。