在CentOS上配置PyTorch的GPU支持涉及几个关键步骤，包括安装NVIDIA驱动、配置CUDA工具包、设置环境变量以及安装PyTorch。以下是详细的步骤指南：
首先，您需要安装适用于您显卡的NVIDIA驱动程序。您可以从NVIDIA官方网站下载并安装驱动程序。
ELRepo是一个第三方软件仓库，提供了许多额外的软件包，包括NVIDIA驱动和CUDA工具包。
# 下载ELRepo仓库的GPG密钥
rpm --import https://www.elrepo.org/RPM-GPG-KEY-elrepo.org
# 安装ELRepo仓库
rpm -Uvh http://www.elrepo.org/elrepo-release-7.0-2.el7.elrepo.noarch.rpm
# 安装nvidia软件包
yum install kmod-nvidia
在安装了ELRepo仓库之后，您可以安装CUDA工具包。请根据您的CUDA版本选择相应的安装命令。例如，如果您需要CUDA 11.1，可以使用以下命令：
# 更新yum源
yum clean all
yum makecache
yum -y update
# 安装CUDA
wget https://developer.download.nvidia.com/compute/cuda/11.1/Prod/local_installers/cuda-repo-rhel7-11-1-local-11.1.85-450.51.06-1.el7.x86_64.rpm
sudo rpm -i cuda-repo-rhel7-11-1-local-11.1.85-450.51.06-1.el7.x86_64.rpm
sudo yum clean all
sudo yum -y install cuda
cuDNN是用于深度神经网络的GPU加速库。您需要根据您的CUDA版本安装相应的cuDNN版本。例如，对于CUDA 11.1，您可以从NVIDIA官方网站下载并安装cuDNN 8.x版本。
nvidia-docker是一个允许在Docker容器中使用NVIDIA GPU的工具。
# 安装Docker
sudo yum install docker
# 启用NVIDIA容器运行时仓库
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.repo | sudo tee /etc/yum.repos.d/nvidia-docker.repo
# 安装nvidia-docker2
sudo yum install -y nvidia-docker2
# 重启Docker服务
sudo systemctl restart docker
安装完成后，您可以使用以下命令验证安装是否成功：
# 拉取nvidia的镜像
docker pull nvidia/cuda:11.1.1-cudnn8-runtime-ubuntu18.04
# 运行一个Docker容器并验证nvidia-smi
docker run --gpus all nvidia/cuda:11.1.1-cudnn8-runtime-ubuntu18.04 nvidia-smi
最后，您可以使用pip在Conda环境中安装PyTorch。请根据您的CUDA版本选择相应的PyTorch版本。例如，对于CUDA 11.1，您可以使用以下命令：
# 创建并激活一个新的Conda环境
conda create -n torch_env python=3.8
conda activate torch_env
# 安装PyTorch
pip install torch torchvision torchaudio -f https://download.pytorch.org/whl/cu111/torch_stable.html
安装完成后，您可以使用以下代码验证PyTorch是否能够使用GPU：
import torch
print(torch.cuda.is_available()) # 应该输出True
print(torch.cuda.device_count()) # 输出可用的GPU数量
通过以上步骤，您应该能够在CentOS上成功配置PyTorch的GPU支持。如果您遇到任何问题，请参考PyTorch官方文档或相关社区论坛寻求帮助。