在CentOS上使用GPU加速PyTorch_运维文库_资讯中心

发布时间:2026-04-27 07:45:43

阅读量:2

在CentOS上使用GPU加速PyTorch，你需要确保你的系统满足以下条件：

一台配备NVIDIA GPU的服务器。
NVIDIA驱动程序已正确安装。
CUDA Toolkit已安装。
cuDNN库已安装。
PyTorch已安装，并且支持CUDA。

以下是详细步骤：

1. 安装NVIDIA驱动程序

首先，你需要安装适合你GPU型号的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

2. 安装CUDA Toolkit

访问NVIDIA CUDA Toolkit下载页面，选择适合你系统的版本并下载。对于CentOS，你可以使用以下命令安装：

# 添加CUDA仓库
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo tee /etc/apt/trusted.gpg.d/nvidia-docker.gpg
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

配置环境变量：

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

验证CUDA Toolkit是否安装成功：

nvcc --version

3. 安装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*

4. 安装PyTorch

使用pip安装支持CUDA的PyTorch版本。你可以从PyTorch官方网站获取安装命令。例如：

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

验证PyTorch是否支持CUDA：

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

如果输出为True，则表示PyTorch已成功配置并可以使用GPU加速。

5. 使用GPU加速PyTorch

在你的PyTorch代码中，你可以使用GPU来加速计算。以下是一个简单的示例：

import torch
import torch.nn as nn
import torch.optim as optim

# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 创建一个张量并将其移动到GPU
x = torch.randn(10, 10).to(device)

# 创建一个模型并将其移动到GPU
model = nn.Linear(10, 10).to(device)

# 前向传播
y = model(x)

print(y)

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

以上就是关于“在CentOS上使用GPU加速PyTorch”的相关介绍，筋斗云是国内较早的云主机应用的服务商，拥有10余年行业经验，提供丰富的云服务器、租用服务器等相关产品服务。云服务器资源弹性伸缩，主机vCPU、内存性能强悍、超高I/O速度、故障秒级恢复；电子化备案，提交快速，专业团队7×24小时服务支持！

简单好用、高性价比云服务器租用链接：https://www.jindouyun.cn/product/cvm