Linux上PyTorch与CUDA如何配合

Linux上 PyTorch 与 CUDA 配合的完整指南

一 核心原理与版本匹配

• 运行 PyTorch 的 GPU 加速需要同时满足：有支持 CUDA 的 NVIDIA 显卡、正确安装的 NVIDIA 驱动、与 PyTorch 构建对应的 CUDA/cuDNN 运行时，以及与之匹配的 PyTorch 预编译包。
• 版本匹配要点：
  • 先用 nvidia-smi 查看驱动与“支持的 CUDA Runtime 版本”（右上角显示，如：12.4）。注意这不是你系统里安装的 Toolkit 版本，但 PyTorch 预编译包通常需要不大于该版本的 CUDA 运行时。
  • 选择与之匹配的 PyTorch 预编译包（pip 的 cu118/cu121 等标签，或 conda 的 cudatoolkit=11.8/12.1 等）。若 torch.cuda.is_available() 为 False，90% 是版本不匹配或装了 CPU 版。
  • 常见误区：驱动太旧、安装了 CPU 版 PyTorch、CUDA/cuDNN 与 PyTorch 版本不对应、环境变量未生效。

二 两种常用部署方式与命令

• 方式 A：系统级 CUDA + cuDNN（适合需要多框架共享 CUDA 的场景）

  1. 安装驱动（Ubuntu 示例）：
     • 查看显卡：lspci | grep -i nvidia
     • 添加显卡驱动 PPA 并安装（示例）：sudo add-apt-repository ppa:graphics-drivers/ppa && sudo apt update && sudo apt install nvidia-driver-535
     • 重启：sudo reboot
  2. 安装 CUDA Toolkit（示例）：
     • wget https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda_11.7.1_515.43.04_linux.run
     • sudo sh cuda_11.7.1_515.43.04_linux.run（若已单独装驱动，安装时可取消驱动选项）
  3. 安装 cuDNN（与 CUDA 版本严格对应）：
     • 解压后复制头文件和库到 CUDA 目录并赋权： tar -xzvf cudnn-11.7-linux-x64-v8.4.1.50.tgz sudo cp cuda/include/cudnn*.h /usr/local/cuda/include sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64 sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*
  4. 配置环境变量（~/.bashrc 或 /etc/profile）：
     • echo ‘export PATH=/usr/local/cuda/bin:$PATH’ >> ~/.bashrc
     • echo ‘export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH’ >> ~/.bashrc
     • source ~/.bashrc
  5. 安装 PyTorch（pip 示例，选择与系统 CUDA 匹配的标签）：
     • pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
  6. 验证：
     • python - <<‘PY’ import torch print(“CUDA可用性:”, torch.cuda.is_available()) print(“GPU数量:”, torch.cuda.device_count()) print(“GPU名称:”, torch.cuda.get_device_name(0)) PY

• 方式 B：conda 环境内嵌 cudatoolkit（更省心，推荐）

  1. 创建并激活环境：conda create -n torch-gpu python=3.10 -y && conda activate torch-gpu
  2. 安装与 PyTorch 匹配的 cudatoolkit 与 PyTorch（示例为 CUDA 11.7）：
     • conda install pytorch torchvision torchaudio cudatoolkit=11.7 -c pytorch
  3. 验证：
     • python - <<‘PY’ import torch print(“CUDA可用性:”, torch.cuda.is_available()) print(“GPU数量:”, torch.cuda.device_count()) print(“GPU名称:”, torch.cuda.get_device_name(0)) PY 说明：conda 会将 cudatoolkit 与 PyTorch 一起安装在环境内，通常无需在系统层面单独配置 CUDA/cuDNN，也更易复现实验环境。

三 在代码中使用 GPU

• 单卡最简用法：
  • device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
  • model.to(device); x = x.to(device); y = model(x)
• 多卡训练建议：
  • 单机多卡：使用 DistributedDataParallel（DDP），比 DataParallel 更高效、可扩展。
  • 多机多卡：同样使用 DDP，配合启动器（如 torchrun/SLURM）进行进程组初始化与通信。

四 常见问题排查清单

• 驱动与兼容性：用 nvidia-smi 确认驱动正常；若 PyTorch 报 “no CUDA-capable device”，多半是驱动过旧或显卡不被支持。
• 版本不匹配：
  • 现象：import torch 正常，但 torch.cuda.is_available() 为 False。
  • 处理：核对 PyTorch 预编译包的 CUDA 标签（如 cu118/cu121）与系统 CUDA/cudatoolkit 是否匹配；必要时重装对应版本。
• 安装了 CPU 版：
  • 现象：torch.cuda.is_available() 恒为 False。
  • 处理：pip 安装时确保使用带 cu 标签的索引；conda 安装时确保包含 cudatoolkit 且不是 cpuonly 包。可用 conda list | grep pytorch 检查构建信息。
• 环境变量未生效：
  • 现象：库找不到、程序启动报 libcuda.so 等缺失。
  • 处理：确认 PATH 与 LD_LIBRARY_PATH 包含 CUDA 的 bin 与 lib64 路径，并执行 source 使其生效。
• 多用户/集群环境：
  • 建议优先使用 conda 内嵌 cudatoolkit 的方案，减少系统级依赖冲突；必要时联系管理员统一驱动与基础环境。