在Ubuntu上為PyTorch設置GPU加速,你需要確保你的系統滿足以下條件:
以下是設置步驟:
首先,你需要安裝適合你GPU型號的NVIDIA驅動程序。你可以從NVIDIA官方網站下載最新的驅動程序。安裝完成后,重啟計算機。
訪問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
使更改生效。
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版本號。
安裝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
安裝完成后,你可以通過運行以下命令來驗證PyTorch是否能夠檢測到GPU:
import torch
print(torch.cuda.is_available())
如果輸出為True
,則表示PyTorch已經成功配置了GPU加速。
請注意,上述步驟可能會隨著CUDA、cuDNN和PyTorch版本的更新而變化。建議查看官方文檔以獲取最新信息。