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 版本，减少系统层面冲突

最新问答