Ubuntu系统中PyTorch的版本兼容性测试 - 问答

Ubuntu 下 PyTorch 版本兼容性测试

一目标与适用范围

面向 Ubuntu 20.04/22.04/24.04，验证 PyTorch 与 CUDA/NVIDIA 驱动/显卡算力/Python 的匹配性，覆盖 pip 与 conda 两种安装方式，确保 GPU 可正常调用并满足项目所需版本组合。
关键关注点：PyTorch 实际绑定的 CUDA 运行时、驱动支持的 最高 CUDA 版本、显卡 SM 算力是否被当前 PyTorch 构建支持、以及 torchvision/torchaudio 的版本对应关系。

二快速检查清单

基础版本信息
- 查看 PyTorch、Torchvision、PyTorch 绑定的 CUDA：
  - python -c “import torch, torchvision; print(‘torch:’, torch.version, ‘torchvision:’, torchvision.version, ‘torch cuda:’, torch.version.cuda)”
- 查看驱动与驱动支持的最高 CUDA（右上角 CUDA Version）：
  - nvidia-smi
- 查看系统 Toolkit（nvcc）版本：
  - nvcc --version
可用性验证
- GPU 是否可用：
  - python -c “import torch; print(‘cuda available:’, torch.cuda.is_available())”
- 设备名称与数量：
  - python -c “import torch; print(‘device count:’, torch.cuda.device_count()); print(‘name:’, torch.cuda.get_device_name(0) if torch.cuda.is_available() else ‘N/A’)”
- 最小计算测试（避免仅有驱动而无实际计算库）：
  - python - <<‘PY’ import torch if torch.cuda.is_available(): x = torch.randn(2, 2, device=‘cuda’) y = x @ x print(“cuda matmul ok, result:”, y[0,0].item()) else: print(“cuda not available”) PY
算力兼容性（避免“SM 算力不支持”）
- 查看显卡算力：
  - python - <<‘PY’ import torch if torch.cuda.is_available(): print(“GPU arch:”, torch.cuda.get_device_capability()) else: print(“cuda not available”) PY
- 若报错提示不支持的 SM（如 sm_86 未被当前 PyTorch 构建支持），需更换包含对应算力的 PyTorch CUDA 构建或升级 PyTorch 版本。

三判定规则与常见不兼容场景

驱动与运行时关系
- nvidia-smi 右上角的 CUDA Version 表示驱动可支持的最高 CUDA 运行时；PyTorch 可运行的 CUDA 运行时需 ≤ 该值。PyTorch 常随包携带自己的 CUDA 运行时，因此与系统全局 nvcc 版本可能不同，这并不影响运行。
版本对应关系
- 需确保 PyTorch 版本 ↔ 其内置 CUDA 版本匹配；同时 torchvision/torchaudio 与主版本保持一致（如 2.6.0 对应 0.17.0）。
算力匹配
- 新卡（如 RTX 3090/4090，SM 86）需要包含该算力的 PyTorch CUDA 构建；否则会出现类似 “not compatible with the current PyTorch installation” 的提示。
典型不兼容与修复
- 仅 CPU 包却期望 GPU：安装带 cu118/cu121/cu126 等标签的 GPU 版本。
- 驱动过旧：升级驱动，使 nvidia-smi 的 CUDA Version ≥ PyTorch 的 torch.version.cuda。
- 算力不支持：更换为包含目标 SM 的 PyTorch 构建或升级 PyTorch。

四一键测试脚本与结果解读

将以下脚本保存为 test_torch_compat.sh，赋予可执行权限并运行：
- chmod +x test_torch_compat.sh && ./test_torch_compat.sh

#!/usr/bin/env bash
set -e

echo "=== PyTorch 兼容性快速测试 ==="

# 0) 环境信息
echo -e "\n[0] 环境信息"
python3 - <<'PY'
import sys, torch, torchvision, platform, subprocess, re
print("Python:", sys.version)
print("Platform:", platform.platform())
print("torch:", torch.__version__, "torchvision:", torchvision.__version__, "torch cuda:", torch.version.cuda)
print("torch cuda available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("device count:", torch.cuda.device_count())
    print("device name:", torch.cuda.get_device_name(0))
    print("GPU arch:", torch.cuda.get_device_capability())
PY

# 1) 驱动与 nvcc
echo -e "\n[1] 驱动与 nvcc"
if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi | head -n 4 | tail -n 2
else
  echo "nvidia-smi not found"
fi
if command -v nvcc >/dev/null 2>&1; then
  nvcc --version | head -n 4 | tail -n 1
else
  echo "nvcc not found"
fi

# 2) 最小计算测试
echo -e "\n[2] 最小计算测试"
python3 - <<'PY'
import torch
if torch.cuda.is_available():
    x = torch.randn(2, 2, device='cuda')
    y = x @ x
    z = y.sum()
    print("cuda matmul & sum ok, result:", z.item())
else:
    print("cuda not available")
PY

# 3) 判定提示
echo -e "\n[3] 判定提示"
echo "- 若 'torch.cuda.is_available()' 为 False："
echo "  1) 可能安装了 CPU 版 PyTorch；2) 驱动未安装/过旧；3) 驱动支持的 CUDA 低于 PyTorch 的 CUDA 运行时。"
echo "- 若报 'SM 算力不支持'：请安装包含该 GPU 算力的 PyTorch CUDA 构建或升级 PyTorch 版本。"
echo "- 三者关系要点："
echo "  nvidia-smi 的 'CUDA Version' >= torch.version.cuda；"
echo "  PyTorch 可自带 CUDA 运行时，和 nvcc 版本可不同；"
echo "  torchvision/torchaudio 版本需与主版本匹配。"

结果解读要点
- 全部通过：环境满足目标 PyTorch + CUDA + 驱动 + 算力 的兼容性要求。
- 仅 CPU 可用：检查是否误装 CPU 版，或 驱动/CUDA 不匹配，必要时重装带 cu118/cu121/cu126 的 GPU 版本。
- 算力不支持：更换为包含目标 SM 的构建或升级 PyTorch。

五复现与修复命令示例

安装指定 CUDA 构建（pip）
- CUDA 12.1：pip install torch==2.1.0+cu121 torchvision==0.16.0+cu121 torchaudio==2.1.0+cu121 --index-url https://download.pytorch.org/whl/cu121
- 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
- CPU 版：pip install torch==2.6.0 torchvision==0.17.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cpu
安装指定 CUDA 构建（conda）
- conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
驱动建议
- Ubuntu 推荐安装与所用 CUDA 运行时匹配的驱动；可使用 ubuntu-drivers autoinstall 自动安装合适版本，或指定版本如：sudo apt install nvidia-driver-555（示例）。

0 赞

0 踩