溫馨提示×

Ubuntu上PyTorch的GPU加速如何設置

小樊
50
2025-09-25 14:23:16
欄目: 智能運維

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

  1. 一臺配備了NVIDIA GPU的計算機。
  2. 已安裝正確版本的NVIDIA驅動程序。
  3. CUDA Toolkit已安裝。
  4. cuDNN庫已安裝。
  5. PyTorch已安裝,并且是為CUDA版本編譯的。

以下是設置步驟:

步驟 1: 安裝NVIDIA驅動程序

首先,你需要安裝適合你GPU型號的NVIDIA驅動程序。你可以從NVIDIA官方網站下載最新的驅動程序。安裝完成后,重啟計算機。

步驟 2: 安裝CUDA Toolkit

訪問NVIDIA的CUDA Toolkit下載頁面,選擇適合你的操作系統和GPU型號的CUDA版本。下載并安裝CUDA Toolkit。安裝完成后,你需要將CUDA的路徑添加到環境變量中??梢酝ㄟ^編輯~/.bashrc文件來實現:

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

然后運行source ~/.bashrc使更改生效。

步驟 3: 安裝cuDNN

cuDNN是NVIDIA的深度學習庫,它提供了許多優化的GPU加速函數。你需要注冊NVIDIA開發者賬戶并下載與你的CUDA版本兼容的cuDNN庫。下載后,解壓縮并將文件復制到CUDA目錄中:

tar -xzvf cudnn-11.x-linux-x64-v8.x.x.x.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*

請確保將x.x替換為你下載的cuDNN版本號。

步驟 4: 安裝PyTorch

安裝PyTorch時,你需要選擇一個與你的CUDA版本兼容的PyTorch版本。你可以使用pip或conda來安裝。例如,如果你想安裝支持CUDA 11.3的PyTorch,可以使用以下命令:

使用pip安裝:

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

或者使用conda安裝:

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

步驟 5: 驗證安裝

安裝完成后,你可以通過運行以下命令來驗證PyTorch是否能夠檢測到GPU:

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

如果輸出為True,則表示PyTorch已經成功配置了GPU加速。

請注意,上述步驟可能會隨著CUDA、cuDNN和PyTorch版本的更新而變化。建議查看官方文檔以獲取最新信息。

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