Ubuntu系统中PyTorch版本如何选择

Ubuntu下PyTorch版本选择指南

一 决策顺序与优先级

• 明确硬件与用途：是否有 NVIDIA GPU（决定 CPU 或 GPU 版本）、是否需要 CUDA/cuDNN、是否需要 多卡/分布式、是否需要 ROCm（AMD GPU）。
• 对齐生态版本：深度学习框架（如 Transformers、Detectron2、MMDetection）与 Python 版本要求，优先选择与之兼容的 PyTorch 稳定版本。
• 选择安装方式：优先使用 虚拟环境（conda 或 venv），避免与系统包冲突；GPU 场景建议用 conda 管理 cudatoolkit，减少与系统 CUDA 的耦合。
• 保持可复现：固定 PyTorch、torchvision、torchaudio 的版本三元组，并记录安装命令与平台标签（如 cu121、cpu）。

二 关键依赖与兼容性核对

• 检查系统 CUDA 与驱动：
  • 查看 CUDA 编译器版本：nvcc --version
  • 查看驱动是否到位：能正常加载 nvidia 内核模块且版本满足目标 CUDA
• 检查 Python 版本：建议 Python ≥ 3.8；创建虚拟环境后再安装 PyTorch
• 核对 PyTorch 与 CUDA 的对应关系：每个 PyTorch 版本只支持特定的 CUDA/cuDNN 组合，安装前以官方兼容表为准
• 验证安装是否成功：
  • 在 Python 中执行：
    • import torch
    • print(“PyTorch version:”, torch.version)
    • print(“CUDA is available:”, torch.cuda.is_available())
  • 若 torch.cuda.is_available() 为 False，通常是 CUDA 不匹配或驱动问题

三 常见场景与推荐选择

• 无 NVIDIA GPU 或仅做 CPU 推理：
  • 选择 CPU 版本，安装命令示例：
    • pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
• 有 NVIDIA GPU（优先使用 conda 管理 cudatoolkit）：
  • 若系统 CUDA 为 11.7：
    • conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
  • 若系统 CUDA 为 11.3：
    • conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
• 需要指定 CUDA 小版本（pip 直装预编译包）：
  • 例如 CUDA 12.1：
    • pip3 install torch==2.1.0+cu121 torchvision==0.16.0+cu121 torchaudio==2.1.0+cu121 --index-url https://download.pytorch.org/whl/cu121
• 需要固定版本三元组（便于复现）：
  • 示例（CPU）：pip3 install torch==2.6.0 torchvision==0.17.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cpu
  • 示例（GPU）：pip3 install torch==2.6.0 torchvision==0.17.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu126
• 使用 Poetry 管理多平台依赖（示例思路）：
  • 为 pytorch-cpu 与 pytorch-cu126 分别定义源，按 sys_platform 与 platform_machine 选择对应版本与源，便于在 Linux x86_64 与 CPU/GPU 场景间切换

四 不匹配与升级降级的处理

• 快速诊断：
  • 在 Python 中查看：torch.version、torch.version.cuda、torch.cuda.is_available()
  • 若 CUDA 不可用或版本不符，优先检查驱动与 CUDA 安装，再核对 PyTorch 的 CUDA 标签是否匹配
• 升级或降级 CUDA（示例思路）：
  • 升级：从 NVIDIA 官网获取目标版本 CUDA Toolkit 的 .deb 包，执行安装（如 sudo apt-get install cuda-）
  • 降级：先清理旧版本（如 sudo apt-get --purge remove “cuda”），再安装目标版本
  • 注意：升级/降级后需重启并重新验证 torch.cuda.is_available()
• 更稳妥的 GPU 环境策略：
  • 使用 conda 安装 cudatoolkit（不改动系统 CUDA），选择与项目匹配的 cudatoolkit 版本，减少系统层面冲突