CentOS上如何安装PyTorch GPU版本

在CentOS上安装PyTorch GPU版本需要经过几个步骤，包括安装NVIDIA驱动、CUDA Toolkit和cuDNN库。以下是详细的步骤：

1. 安装NVIDIA驱动

首先，确保你的系统已经安装了NVIDIA显卡，并且已经下载了相应的驱动程序。

检查显卡型号

lspci | grep -i nvidia

下载并安装驱动

1. 访问NVIDIA驱动下载页面。
2. 选择你的显卡型号和操作系统版本，下载相应的驱动程序。
3. 安装驱动程序：

   sudo sh NVIDIA-Linux-x86_64-xxx.xx.run
   

   按照提示完成安装。

2. 安装CUDA Toolkit

下载CUDA Toolkit

1. 访问CUDA Toolkit下载页面。
2. 选择你的操作系统和版本，下载相应的CUDA Toolkit。

安装CUDA Toolkit

1. 解压下载的文件：

   tar -xzvf cuda_xxx.x_linux.run
   

2. 进入解压后的目录并运行安装脚本：

   cd cuda_xxx.x_linux
   sudo sh cuda_xxx.x_linux.run
   

3. 按照提示完成安装。在安装过程中，确保选择了安装CUDA Toolkit和驱动程序。

配置环境变量

编辑~/.bashrc文件，添加以下内容：

export PATH=/usr/local/cuda-x.x/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-x.x/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

然后运行：

source ~/.bashrc

3. 安装cuDNN库

下载cuDNN

1. 访问cuDNN下载页面。
2. 注册并登录，下载与你的CUDA版本兼容的cuDNN库。

安装cuDNN

1. 解压下载的文件：

   tar -xzvf cudnn-x.x-linux-x64-vxxx.tgz
   

2. 将解压后的文件复制到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 GPU版本

使用pip安装

你可以使用pip来安装PyTorch GPU版本。首先，确保你已经安装了pip：

sudo yum install python-pip

然后，安装PyTorch GPU版本：

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

注意：cu113表示CUDA 11.3版本，你需要根据你安装的CUDA版本选择相应的URL。

5. 验证安装

最后，验证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可用并且显示了GPU型号，说明安装成功。

通过以上步骤，你应该能够在CentOS上成功安装PyTorch GPU版本。