PyTorch在CentOS上的GPU加速如何设置

在CentOS上设置PyTorch的GPU加速，需按以下步骤操作：

1. 安装NVIDIA驱动

   • 使用命令检查驱动是否已安装：nvidia-smi。
   • 若未安装，从NVIDIA官网下载对应型号的驱动（如.run文件），执行安装并重启系统。
   • 安装时需禁用nouveau驱动（通过修改/etc/default/grub和/lib/modprobe.d/dist-blacklist.conf）。

2. 安装CUDA Toolkit

   • 从NVIDIA官网下载与驱动兼容的CUDA版本（如CUDA 11.8），通过rpm或run文件安装。
   • 添加环境变量：

     echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc  
     echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc  
     source ~/.bashrc  
     

3. 安装cuDNN

   • 从NVIDIA官网下载与CUDA版本匹配的cuDNN库，解压后复制文件至CUDA目录：

     tar -xzvf cudnn-*.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. 安装支持GPU的PyTorch

   • 使用pip安装时，指定CUDA版本（如CUDA 11.8）：

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

   • 或通过conda安装：

     conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia  
     

5. 验证安装

   • 运行Python代码检查GPU是否可用：

     import torch  
     print(torch.cuda.is_available())  # 应返回True  
     print(torch.cuda.get_device_name(0))  # 显示GPU型号  
     

6. 使用GPU加速

   • 在代码中将模型和数据移动到GPU：

     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
     model.to(device)  
     data = data.to(device)  
     

注意：

• 需确保CUDA、cuDNN与PyTorch版本兼容，可参考PyTorch官网的版本对照表。
• 多GPU环境可使用torch.nn.DataParallel或DistributedDataParallel。