要在CentOS上使用PyTorch进行GPU加速,您需要完成以下几个步骤:
确保您的CentOS系统上安装了支持CUDA的NVIDIA GPU。
首先,您需要安装NVIDIA显卡驱动。可以参考NVIDIA官方文档进行安装。
# 添加EPEL仓库
sudo yum install epel-release
# 安装NVIDIA驱动
sudo yum install kernel-devel-$(uname -r) kernel-headers-$(uname -r)
sudo yum install nvidia-driver-latest-dkms
# 重启系统
sudo reboot
验证驱动是否安装成功:
nvidia-smi
访问NVIDIA CUDA Toolkit下载页面,选择适合您系统的版本并下载。然后按照官方指南进行安装。
例如,使用yum安装CUDA Toolkit:
# 添加CUDA仓库
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/yum.repos.d/nvidia-docker.repo
# 安装CUDA Toolkit
sudo yum install cuda
访问NVIDIA cuDNN下载页面,下载与您的CUDA版本兼容的cuDNN库。然后解压并复制到CUDA目录。
# 下载cuDNN
tar -xzvf cudnn-11.7-linux-x64-v8.6.0.26.tgz
# 复制文件到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*
如果您需要进行分布式训练,还需要安装NCCL。
# 下载NCCL
tar -xzvf nccl_2.15.3-1+cuda11.7_x86_64.txz
# 复制文件到CUDA目录
sudo cp nccl/include/nccl.h /usr/local/cuda/include
sudo cp -P nccl/lib/libnccl* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/nccl.h /usr/local/cuda/lib64/libnccl*
使用pip安装PyTorch,确保选择与您的CUDA版本兼容的PyTorch版本。
# 创建一个新的虚拟环境(可选)
python3 -m venv pytorch_env
source pytorch_env/bin/activate
# 安装PyTorch
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
验证PyTorch是否能够检测到GPU。
import torch
print(torch.cuda.is_available())
print(torch.cuda.current_device())
print(torch.cuda.get_device_name(0))
如果输出显示GPU设备信息,说明PyTorch已经成功配置并可以使用GPU加速。
为了方便使用CUDA和cuDNN,可以配置环境变量。
编辑~/.bashrc文件,添加以下内容:
export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
然后使配置生效:
source ~/.bashrc
通过以上步骤,您应该能够在CentOS上成功配置并使用PyTorch进行GPU加速。