CentOS上PyTorch与CUDA如何配合

centos

小樊

2025-03-01 19:17:18

栏目: 智能运维

在CentOS上配置PyTorch与CUDA的配合，可以按照以下步骤进行：

1. 安装NVIDIA驱动

首先，确保你的系统已经安装了NVIDIA显卡驱动。你可以通过以下命令检查是否已经安装了驱动：

nvidia-smi

如果没有安装，可以参考NVIDIA官方文档进行安装。

2. 安装CUDA Toolkit

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

sudo yum install https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-repo-rhel7-10.2.89-1.x86_64.rpm
sudo yum clean all
sudo yum install cuda

3. 配置环境变量

编辑~/.bashrc文件，添加以下内容以配置CUDA路径：

export PATH=/usr/local/cuda-10.2/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-10.2/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

然后使配置生效：

source ~/.bashrc

4. 安装cuDNN

访问NVIDIA cuDNN下载页面，选择与你的CUDA版本兼容的cuDNN版本并下载。对于CentOS，可以使用以下命令安装：

tar -xzvf cudnn-10.2-linux-x64-v8.0.5.39.tgz
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. 安装PyTorch

你可以使用pip来安装PyTorch。首先，确保你已经安装了pip：

sudo yum install python3-pip

然后，根据你的CUDA版本选择合适的PyTorch安装命令。例如，如果你使用的是CUDA 10.2，可以使用以下命令：

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

6. 验证安装

最后，验证PyTorch是否能够检测到CUDA：

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

如果输出为True和10.2，则表示PyTorch已经成功配置并与CUDA配合使用。

注意事项

确保你的系统满足所有依赖项要求。
如果遇到问题，可以参考PyTorch官方文档和社区论坛寻求帮助。

通过以上步骤，你应该能够在CentOS上成功配置PyTorch与CUDA的配合。

亿速云「云服务器」，即开即用、新一代英特尔至强铂金CPU、三副本存储NVMe SSD云盘，价格低至29元/月。点击查看>>