CentOS下如何使用GPU加速PyTorch

在CentOS系统上使用GPU加速PyTorch，你需要完成以下几个步骤：

1. 安装NVIDIA驱动： 首先，确保你的CentOS系统上安装了与你的GPU兼容的NVIDIA驱动。你可以从NVIDIA官网下载对应的驱动程序。

   sudo yum install epel-release
   sudo yum install dkms
   sudo yum install nvidia-driver-latest-dkms
   

   安装完成后，重启系统并验证驱动是否安装成功：

   nvidia-smi
   

   如果系统显示了GPU的信息，说明驱动安装成功。

2. 安装CUDA Toolkit： PyTorch需要CUDA Toolkit来支持GPU加速。你可以从NVIDIA官网下载适合你系统的CUDA Toolkit版本。

   # 下载CUDA Toolkit
   wget https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-repo-rhel7-10.2.89-1.x86_64.rpm
   sudo yum localinstall cuda-repo-rhel7-10.2.89-1.x86_64.rpm
   sudo yum clean all
   sudo yum install cuda
   
   # 设置环境变量
   echo 'export PATH=/usr/local/cuda-10.2/bin:$PATH' >> ~/.bashrc
   echo 'export LD_LIBRARY_PATH=/usr/local/cuda-10.2/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
   source ~/.bashrc
   

   验证CUDA是否安装成功：

   nvcc --version
   

3. 安装cuDNN： cuDNN是用于深度神经网络的GPU加速库。你需要注册NVIDIA开发者账号并下载对应的cuDNN库。

   下载cuDNN后，解压并将文件复制到CUDA目录：

   tar -xzvf cudnn-10.2-linux-x64-v8.0.5.39.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. 安装NCCL（可选）： 如果你需要进行多GPU训练或者分布式训练，你可能还需要安装NCCL（NVIDIA Collective Communications Library）。

   下载并安装NCCL：

   # 下载NCCL
   wget https://developer.download.nvidia.com/compute/machine-learning/repos/rhel7/x86_64/libnccl2-2.7.7-1+cuda10.2.x86_64.rpm
   wget https://developer.download.nvidia.com/compute/machine-learning/repos/rhel7/x86_64/libnccl-dev-2.7.7-1+cuda10.2.x86_64.rpm
   
   # 安装NCCL
   sudo yum install libnccl2-2.7.7-1+cuda10.2.x86_64.rpm libnccl-dev-2.7.7-1+cuda10.2.x86_64.rpm
   

5. 安装PyTorch： 安装支持CUDA的PyTorch版本。你可以使用pip来安装：

   pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu102
   

   这里的cu102表示PyTorch支持CUDA 10.2。根据你安装的CUDA版本，你可能需要替换为cu110（对应CUDA 11.0）或cu111（对应CUDA 11.1）等。

6. 验证安装： 安装完成后，你可以运行一个简单的PyTorch脚本来验证GPU是否可用：

   import torch
   print(torch.cuda.is_available())
   print(torch.cuda.current_device())
   print(torch.cuda.get_device_name(torch.cuda.current_device()))
   

   如果输出显示了GPU的名称，并且torch.cuda.is_available()返回True，那么PyTorch已经成功配置并可以使用GPU加速了。

请注意，上述步骤可能会随着CUDA、cuDNN和PyTorch版本的更新而变化。建议访问PyTorch官网获取最新的安装指南。