CentOS上PyTorch的GPU加速设置

# 添加EPEL仓库
sudo yum install epel-release

# 安装NVIDIA驱动
sudo yum install nvidia-driver-latest-dkms

# 重启系统
sudo reboot

# 下载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

# 下载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*

# 安装PyTorch和torchvision
pip install torch torchvision torchaudio

# 创建一个新的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

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