温馨提示×

PyTorch在CentOS上如何使用GPU加速

小樊
45
2025-11-16 04:47:52
栏目: 智能运维

要在CentOS上使用PyTorch进行GPU加速,您需要完成以下几个步骤:

1. 硬件准备

确保您的CentOS系统上安装了支持CUDA的NVIDIA GPU。

2. 安装NVIDIA驱动

首先,您需要安装NVIDIA显卡驱动。可以参考NVIDIA官方文档进行安装。

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

# 安装NVIDIA驱动
sudo yum install kernel-devel-$(uname -r) kernel-headers-$(uname -r)
sudo yum install nvidia-driver-latest-dkms

# 重启系统
sudo reboot

验证驱动是否安装成功:

nvidia-smi

3. 安装CUDA Toolkit

访问NVIDIA CUDA Toolkit下载页面,选择适合您系统的版本并下载。然后按照官方指南进行安装。

例如,使用yum安装CUDA Toolkit:

# 添加CUDA仓库
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/yum.repos.d/nvidia-docker.repo

# 安装CUDA Toolkit
sudo yum install cuda

4. 安装cuDNN

访问NVIDIA cuDNN下载页面,下载与您的CUDA版本兼容的cuDNN库。然后解压并复制到CUDA目录。

# 下载cuDNN
tar -xzvf cudnn-11.7-linux-x64-v8.6.0.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*

5. 安装NCCL(可选)

如果您需要进行分布式训练,还需要安装NCCL。

# 下载NCCL
tar -xzvf nccl_2.15.3-1+cuda11.7_x86_64.txz

# 复制文件到CUDA目录
sudo cp nccl/include/nccl.h /usr/local/cuda/include
sudo cp -P nccl/lib/libnccl* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/nccl.h /usr/local/cuda/lib64/libnccl*

6. 安装PyTorch

使用pip安装PyTorch,确保选择与您的CUDA版本兼容的PyTorch版本。

# 创建一个新的虚拟环境(可选)
python3 -m venv pytorch_env
source pytorch_env/bin/activate

# 安装PyTorch
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117

7. 验证安装

验证PyTorch是否能够检测到GPU。

import torch
print(torch.cuda.is_available())
print(torch.cuda.current_device())
print(torch.cuda.get_device_name(0))

如果输出显示GPU设备信息,说明PyTorch已经成功配置并可以使用GPU加速。

8. 配置环境变量(可选)

为了方便使用CUDA和cuDNN,可以配置环境变量。

编辑~/.bashrc文件,添加以下内容:

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

然后使配置生效:

source ~/.bashrc

通过以上步骤,您应该能够在CentOS上成功配置并使用PyTorch进行GPU加速。

0