PyTorch在CentOS上的GPU加速支持

在 CentOS 上启用 PyTorch 的 GPU 加速

支持概览与前置条件

• 在 CentOS 上启用 PyTorch GPU 加速是可行的，前提是系统具备 NVIDIA GPU，并正确安装匹配的 NVIDIA 驱动、CUDA Toolkit、cuDNN，然后安装与之匹配的 PyTorch CUDA 版本。常见流程包括：驱动检查与安装、CUDA 安装与环境变量配置、cuDNN 安装、PyTorch 安装与验证。多机或多卡训练可进一步安装 NCCL 提升通信性能。

标准安装步骤

1. 检查 GPU 与驱动
   • 执行 nvidia-smi 查看驱动与 GPU 状态；若未安装，先从 NVIDIA 官网获取与显卡型号匹配的驱动。
2. 安装 CUDA Toolkit
   • 选择与目标 PyTorch 版本匹配的 CUDA 版本，下载安装包（如 .run 文件），执行安装并按提示完成。
   • 配置环境变量（示例）：
     • export PATH=/usr/local/cuda/bin:$PATH
     • export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
     • 使配置生效：source ~/.bashrc
3. 安装 cuDNN（与 CUDA 版本严格匹配）
   • 下载对应 cuDNN，解压后将头文件与库复制到 CUDA 目录，并设置权限：
     • 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. 安装 PyTorch（pip 或 conda）
   • 选择与已安装 CUDA 版本一致的 PyTorch 预编译包。
5. 验证安装
   • 运行 Python 验证：
     • import torch
     • print(torch.cuda.is_available())
     • print(torch.cuda.current_device())
     • print(torch.cuda.get_device_name(torch.cuda.current_device()))
   • 返回 True 与正确的 GPU 型号 即表示配置成功。

版本匹配与安装命令示例

• 下表给出常见组合与对应安装方式（示例命令，按你的实际版本替换）：

• 说明：cuDNN 必须与 CUDA 版本严格匹配；PyTorch 的 CUDA 标签（如 cu117）需与系统 CUDA 一致。若使用其他 CUDA 版本（如 11.3），请选择对应的 PyTorch 索引或 cudatoolkit 版本。

容器化与多 GPU 场景

• Docker 中使用 GPU：安装 Docker 后，配置 NVIDIA Container Toolkit（nvidia-docker2），即可在容器内使用 nvidia-smi 与 GPU。
• 多 GPU 与分布式训练：安装 NCCL 以优化多卡通信；在代码中使用 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 并将模型与数据 .to(device) 完成迁移。

常见问题与排查

• nvidia-smi 无法执行或找不到命令：驱动未装或内核模块未加载，检查驱动安装日志与 lsmod | grep nvidia。
• torch.cuda.is_available() 返回 False：
  • 核对 CUDA 与 cuDNN 版本是否匹配；
  • 检查环境变量 PATH、LD_LIBRARY_PATH 是否包含 CUDA 的 bin 与 lib64；
  • 确认 PyTorch 安装命令的 CUDA 标签与系统 CUDA 一致；
  • 若使用 conda，避免与系统 CUDA 冲突（优先使用 cudatoolkit 管理）。
• 多卡训练异常：确认 NCCL 安装正确，驱动与 CUDA 版本满足要求，并使用 torch.distributed 的正确启动方式。