在CentOS上为PyTorch设置GPU加速,你需要确保你的系统满足以下条件:
- 一台配备了NVIDIA GPU的计算机。
- 已安装CUDA Toolkit。
- 已安装cuDNN库。
- 安装了支持GPU的PyTorch版本。
以下是详细步骤:
1. 安装NVIDIA驱动
首先,确保你的NVIDIA GPU驱动已经安装并正确配置。
# 添加EPEL仓库 sudo yum install epel-release # 安装NVIDIA驱动 sudo yum install nvidia-driver-latest-dkms # 重启系统 sudo reboot
2. 安装CUDA Toolkit
访问NVIDIA CUDA Toolkit下载页面,选择适合你系统的版本并下载安装脚本。然后运行以下命令:
# 下载CUDA Toolkit安装脚本 wget https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-repo-rhel7-11.4.2-1.x86_64.rpm # 安装CUDA Toolkit sudo yum localinstall cuda-repo-rhel7-11.4.2-1.x86_64.rpm # 安装依赖包 sudo yum install cuda # 配置环境变量 echo 'export PATH=/usr/local/cuda-11.4/bin:$PATH' >> ~/.bashrc echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc # 使环境变量生效 source ~/.bashrc
3. 安装cuDNN库
访问NVIDIA cuDNN下载页面,下载适合你CUDA版本的cuDNN库。然后运行以下命令:
# 下载cuDNN库 tar -xzvf cudnn-11.4-linux-x64-v8.2.2.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*
4. 安装支持GPU的PyTorch版本
你可以使用pip或conda来安装支持GPU的PyTorch版本。以下是使用pip的示例:
# 安装PyTorch和torchvision pip install torch torchvision torchaudio
如果你使用conda,可以运行以下命令:
# 创建一个新的conda环境(可选) conda create -n pytorch_env python=3.9 # 激活环境 conda activate pytorch_env # 安装PyTorch和torchvision conda install pytorch torchvision torchaudio cudatoolkit=11.4 -c pytorch
5. 验证安装
运行以下Python代码来验证PyTorch是否能够检测到GPU:
import torch # 检查是否有可用的GPU if torch.cuda.is_available(): print("GPU is available!") print("Device name:", torch.cuda.get_device_name(0)) else: print("GPU is not available.")
如果一切正常,你应该会看到类似以下的输出:
GPU is available! Device name: NVIDIA GeForce GTX 1080 Ti
这样,你就成功地在CentOS上为PyTorch设置了GPU加速。