CentOS上PyTorch的GPU支持如何配置

在CentOS上配置PyTorch的GPU支持涉及几个关键步骤，包括安装必要的依赖项、CUDA Toolkit、cuDNN库以及PyTorch本身。以下是详细的配置指南：

1. 更新系统

首先，确保你的CentOS系统是最新的：

sudo yum update -y

2. 安装CUDA Toolkit

根据你的CentOS版本和GPU型号，从NVIDIA官方网站下载并安装合适的CUDA Toolkit。例如，对于CUDA 11.1，可以使用以下命令安装：

wget https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-repo-rhel7-11.1.0-1.0.329-1.el7.x86_64.rpms
sudo rpm -ivh cuda-repo-rhel7-11.1.0-1.0.329-1.el7.x86_64.rpms
sudo yum clean all
sudo yum install -y cuda

3. 安装cuDNN库

同样地，从NVIDIA官方网站下载并安装与CUDA版本匹配的cuDNN库。例如，对于CUDA 11.1，可以下载cuDNN 8.0.5：

wget https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.0.5.32/Production/11.1_20211031/cudnn-11.1-linux-x64-v8.0.5.32.tgz
tar -xvf cudnn-11.1-linux-x64-v8.0.5.32.tgz
sudo cp cuda/include/cudnn*.h /usr/local/cuda/includes
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

你可以使用pip来安装PyTorch。为了确保安装的版本与你的CUDA版本兼容，可以指定PyTorch的版本。例如，对于CUDA 11.1，可以使用以下命令安装：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu111

或者，如果你使用Anaconda，可以在激活的虚拟环境中安装：

conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch

5. 验证安装

创建一个名为 test_pytorch.py 的Python脚本，内容如下：

import torch
print(torch.__version__)
print(torch.cuda.is_available())

运行此脚本以验证PyTorch是否已成功安装：

python3 test_pytorch.py

如果输出显示PyTorch的版本号和 True（表示CUDA可用），则表示安装成功。

6. 设置GPU可见性（可选）

如果你有多个GPU，并且想要设置哪些GPU对PyTorch可见，可以使用以下环境变量：

export CUDA_VISIBLE_DEVICES=1,3,2

这将使得GPU 1、3和2对PyTorch可见。

通过以上步骤，你应该能够在CentOS上成功安装并配置PyTorch的GPU支持。如果在安装过程中遇到问题，请检查错误日志并确保所有依赖项都已正确安装。