Ubuntu上PyTorch库依赖如何管理

Ubuntu 上 PyTorch 依赖管理实践

一 环境隔离与基础准备

• 使用虚拟环境隔离依赖：优先选择 venv 或 conda。示例（venv）：
  • 创建并激活：python3 -m venv venv && source venv/bin/activate
  • 升级 pip：pip install -U pip
• 系统级依赖（编译/多媒体等常用库）：
  • sudo apt update && sudo apt install -y build-essential cmake git wget unzip yasm pkg-config libopenblas-dev liblapack-dev libjpeg-dev libpng-dev
• 建议 Python 版本：3.8+，以避免与部分预编译包不兼容。

二 安装 PyTorch 与对应依赖

• 选择安装方式
  • pip：安装官方预编译 wheel，灵活控制版本与 CUDA 版本。
  • conda：同时管理 Python 与二进制依赖（如 cudatoolkit），更省心。
• 常用安装命令（示例）
  • CPU 版本（pip）：pip install torch==2.6.0 torchvision==0.17.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cpu
  • GPU 版本（pip，CUDA 12.6）：pip install torch==2.6.0+cu126 torchvision==0.17.0+cu126 torchaudio==2.6.0+cu126 --index-url https://download.pytorch.org/whl/cu126
  • GPU 版本（conda）：conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch -c nvidia
• 驱动与兼容性要点
  • 使用 GPU 版本时，系统需安装与所选 CUDA 兼容的 NVIDIA 驱动；pip 安装的 PyTorch 自带所需 CUDA 运行时库，但驱动仍需系统层面匹配。
  • 选择与 CUDA 版本匹配的 cuDNN（如使用 conda 的 cudatoolkit，通常无需单独安装 cuDNN）。

三 依赖记录与版本锁定

• requirements.txt（pip）
  • 生成锁定文件：pip freeze > requirements.txt
  • 复现环境：pip install -r requirements.txt
• pip-tools（更可控）
  • 安装：pip install pip-tools
  • 维护：echo "torch==2.6.0" > requirements.in
  • 编译：pip-compile requirements.in -o requirements.txt
  • 同步：pip-sync requirements.txt（精确对齐环境）
• conda 环境
  • 导出：conda env export > environment.yml
  • 复现：conda env create -f environment.yml
• 版本对齐建议
  • 固定 torch、torchvision、torchaudio 的匹配版本（如 2.6.0/0.17.0/2.6.0），避免 API/底层库不兼容。

四 多环境切换与多平台配置

• 使用 conda 管理多环境
  • 创建不同 CUDA 版本环境：conda create -n torch118 python=3.10 pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch -c nvidia
  • 切换：conda activate torch118
• 使用 Poetry 管理复杂依赖与平台标记
  • 示例（在 pyproject.toml 中对不同平台指定源/版本）：
    • torch = [{version = "2.6.0", source = "pytorch-cpu", markers = "platform_machine == 'x86_64'"}]
  • 优势：同一仓库支持 CPU/GPU、x86_64/ARM 等多平台。

五 验证与常见问题处理

• 安装验证
  • python - <<'PY' import torch, sys print("torch:", torch.__version__, "cuda:", torch.cuda.is_available()) print("python:", sys.version) PY
• 常见问题
  • 版本冲突：优先使用虚拟环境；用 pip check 排查；必要时 pip install -U 包名 或回退版本。
  • 下载慢/不稳定：pip 可使用国内镜像（如清华源）；conda 可配置清华 channel。
  • 驱动/CUDA 不匹配：确认 NVIDIA 驱动版本满足所选 CUDA/cudatoolkit；必要时重装驱动或切换 PyTorch 的 CUDA 版本。