溫馨提示×

CentOS上PyTorch的GPU加速設置

小樊
44
2025-05-23 09:12:44
欄目: 智能運維

在CentOS上為PyTorch設置GPU加速,你需要確保你的系統滿足以下條件:

  1. 一臺配備了NVIDIA GPU的計算機。
  2. 已安裝CUDA Toolkit。
  3. 已安裝cuDNN庫。
  4. 安裝了支持GPU的PyTorch版本。

以下是詳細步驟:

1. 安裝NVIDIA驅動

首先,確保你的NVIDIA GPU驅動已經安裝并正確配置。

# 添加EPEL倉庫
sudo yum install epel-release

# 安裝NVIDIA驅動
sudo yum install nvidia-driver-latest-dkms

# 重啟系統
sudo reboot

2. 安裝CUDA Toolkit

訪問NVIDIA CUDA Toolkit下載頁面,選擇適合你系統的版本并下載安裝腳本。然后運行以下命令:

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

3. 安裝cuDNN庫

訪問NVIDIA cuDNN下載頁面,下載適合你CUDA版本的cuDNN庫。然后運行以下命令:

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

4. 安裝支持GPU的PyTorch版本

你可以使用pip或conda來安裝支持GPU的PyTorch版本。以下是使用pip的示例:

# 安裝PyTorch和torchvision
pip install torch torchvision torchaudio

如果你使用conda,可以運行以下命令:

# 創建一個新的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

5. 驗證安裝

運行以下Python代碼來驗證PyTorch是否能夠檢測到GPU:

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

這樣,你就成功地在CentOS上為PyTorch設置了GPU加速。

0
亚洲午夜精品一区二区_中文无码日韩欧免_久久香蕉精品视频_欧美主播一区二区三区美女