温馨提示×

CentOS上如何安装PyTorch GPU版本

小樊
47
2025-07-17 01:30:32
栏目: 智能运维

在CentOS上安装PyTorch GPU版本需要经过几个步骤,包括安装NVIDIA驱动、CUDA Toolkit和cuDNN库。以下是详细的步骤:

1. 安装NVIDIA驱动

首先,确保你的系统已经安装了NVIDIA显卡,并且已经下载了相应的驱动程序。

检查显卡型号

lspci | grep -i nvidia 

下载并安装驱动

  1. 访问NVIDIA驱动下载页面
  2. 选择你的显卡型号和操作系统版本,下载相应的驱动程序。
  3. 安装驱动程序:
    sudo sh NVIDIA-Linux-x86_64-xxx.xx.run 
    按照提示完成安装。

2. 安装CUDA Toolkit

下载CUDA Toolkit

  1. 访问CUDA Toolkit下载页面
  2. 选择你的操作系统和版本,下载相应的CUDA Toolkit。

安装CUDA Toolkit

  1. 解压下载的文件:
    tar -xzvf cuda_xxx.x_linux.run 
  2. 进入解压后的目录并运行安装脚本:
    cd cuda_xxx.x_linux sudo sh cuda_xxx.x_linux.run 
  3. 按照提示完成安装。在安装过程中,确保选择了安装CUDA Toolkit和驱动程序。

配置环境变量

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

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

然后运行:

source ~/.bashrc 

3. 安装cuDNN库

下载cuDNN

  1. 访问cuDNN下载页面
  2. 注册并登录,下载与你的CUDA版本兼容的cuDNN库。

安装cuDNN

  1. 解压下载的文件:
    tar -xzvf cudnn-x.x-linux-x64-vxxx.tgz 
  2. 将解压后的文件复制到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 GPU版本

使用pip安装

你可以使用pip来安装PyTorch GPU版本。首先,确保你已经安装了pip:

sudo yum install python-pip 

然后,安装PyTorch GPU版本:

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

注意:cu113表示CUDA 11.3版本,你需要根据你安装的CUDA版本选择相应的URL。

5. 验证安装

最后,验证PyTorch是否能够检测到GPU:

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

如果输出显示GPU可用并且显示了GPU型号,说明安装成功。

通过以上步骤,你应该能够在CentOS上成功安装PyTorch GPU版本。

0