在Ubuntu上运行PyTorch需要注意什么

Ubuntu上运行PyTorch的关键注意事项

一 系统与版本匹配

• 选择受支持的 Ubuntu LTS 版本（如 18.04+），并确认系统的 glibc 等基础依赖满足要求。
• 明确 Python 版本需求（建议 3.8–3.11 区间，具体以所用 PyTorch 版本说明为准）。
• 规划 CPU/GPU 路线：无 NVIDIA GPU 可选 CPU 版；有 NVIDIA GPU 需规划 驱动 / CUDA / cuDNN 与 PyTorch 的版本匹配。
• 建议始终使用 虚拟环境（venv 或 conda）隔离依赖，避免项目间冲突。

二 GPU环境部署要点

• 安装与显卡匹配的 NVIDIA 驱动，可用 sudo ubuntu-drivers autoinstall 或安装指定版本（如 sudo apt install nvidia-driver-555），安装后重启。
• 安装 CUDA Toolkit 与 cuDNN，确保二者版本互相兼容，并与后续安装的 PyTorch 版本匹配。
• 设置环境变量（如使用 CUDA 11.x 时）：
  • export PATH=/usr/local/cuda/bin:$PATH
  • export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
• 验证驱动与 CUDA：
  • nvidia-smi（查看驱动与最高可用 CUDA Runtime）
  • nvcc --version（查看已安装的 CUDA Toolkit 编译器版本）

三 安装方式与版本选择

• 常用安装方式：
  • pip：pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu（CPU）
  • conda：conda install pytorch torchvision torchaudio cpuonly -c pytorch（CPU）
  • GPU 版（示例为 CUDA 11.7）：
    • pip：pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
    • conda：conda install pytorch torchvision torchaudio cudatoolkit=11.7 -c pytorch
• 选择依据：优先以 PyTorch 官方安装命令生成器为准，确保所选 PyTorch 版本与系统的 CUDA/cuDNN/驱动兼容。
• 国内网络可临时使用镜像源加速（如清华源），但务必确认包来源可信。

四 验证与常见问题处理

• 安装后验证：
  • python -c "import torch; print(torch.__version__)"
  • python -c "import torch; print(torch.cuda.is_available())"
  • 如为 GPU 环境，进一步查看设备信息：torch.cuda.device_count()、torch.cuda.get_device_name(0)。
• 常见问题与对策：
  • 版本不匹配（PyTorch 与 CUDA/cuDNN/驱动）：按官方兼容矩阵重装对应版本。
  • 网络超时/慢：更换国内镜像源或增大 pip 超时。
  • 驱动/CUDA 不可用：检查 nvidia-smi 与 nvcc --version 输出，确认驱动加载与 Toolkit 安装正确。
  • 权限/依赖缺失：避免在系统 Python 全局安装；必要时使用 --user 或虚拟环境；补齐构建/运行依赖。
  • 环境变量错误：核对 PATH、LD_LIBRARY_PATH、CUDA_HOME 是否指向正确 CUDA 目录。

五 性能与稳定性优化

• 训练加速：使用 混合精度训练（torch.cuda.amp）、高效 DataLoader（多线程/预取/合适 num_workers）、数据解码优化（如 jpeg4py/turbojpeg）。
• 资源监控：使用 nvidia-smi、htop 观察 GPU/CPU/内存占用，及时排查瓶颈。
• 环境与依赖：优先 Anaconda 管理二进制依赖；必要时用 Bottleneck 等工具定位性能瓶颈。