linux

PyTorch在Linux上的版本选择

小樊
45
2025-12-24 10:03:55
栏目: 智能运维

Linux 上选择 PyTorch 版本的关键维度

快速判定与选择步骤

  1. 查看发行版与架构:cat /etc/os-releaseuname -m(常见为 x86_64/AMD64)。
  2. 查看 Python 版本:python -Vpython3 -V
  3. 查看 GPU 与驱动/CUDA:
    • nvidia-smi(右上显示 Supported/Runtime CUDA 版本,表示驱动可支持的最高 CUDA 运行时)
    • nvcc --version(显示已安装的 CUDA Toolkit 版本)
  4. 选择 PyTorch 构建:
    • GPU:选择与驱动/Runtime CUDA 匹配的 PyTorch CUDA 标签(如 cu121 表示 CUDA 12.1 构建)。
    • CPU:直接选 CPU-only 构建。
  5. 选择安装通道与命令:
    • pip:使用官方 whl 索引(如 --extra-index-url https://download.pytorch.org/whl/cu121/cpu)。
    • conda:使用 pytorch-cuda=xxcudatoolkit=xx 指定 CUDA 版本。
  6. 建虚拟环境并安装,完成后验证:import torch; print(torch.__version__, torch.cuda.is_available())

版本匹配与命令示例

场景 判定要点 pip 安装示例 conda 安装示例
CPU-only 无 NVIDIA GPU 或不启用 GPU pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu conda install pytorch torchvision torchaudio cpuonly -c pytorch
GPU + CUDA 12.x nvidia-smi 显示 Runtime ≥ 12.x pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu121 conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
GPU + CUDA 11.8 nvidia-smi 支持 11.8 且已装 11.8 Toolkit pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118 conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch

常见问题与建议

0
看了该问题的人还看了