Ubuntu下PyTorch的GPU支持情況如何

ubuntu下安裝pytorch并利用gpu進行加速計算是深度學習開發中的常見需求。以下是關于在ubuntu下安裝pytorch及其gpu支持情況的詳細說明:

安裝NVIDIA顯卡驅動

首先,需要在Ubuntu上安裝NVIDIA顯卡驅動。可以通過以下命令安裝:

sudo apt update sudo apt install nvidia-driver-版本號 

例如,對于NVIDIA GeForce GTX 1050 Ti顯卡,推薦的驅動版本是470。

安裝CUDA

安裝CUDA是使用PyTorch GPU功能的關鍵步驟。CUDA的版本需要與PyTorch支持的版本相匹配。可以通過以下命令安裝CUDA:

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2004-11-8-local_11.8.0-1_amd64.deb sudo dpkg -i cuda-repo-ubuntu2004-11-8-local_11.8.0-1_amd64.deb sudo cp /var/cuda-repo-ubuntu2004-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/ sudo apt-get update sudo apt-get -y install cuda 

安裝cuDNN

接下來,需要安裝與CUDA版本相匹配的cuDNN庫:

sudo apt install libcudnn8 

安裝PyTorch GPU版本

最后,使用condapip安裝PyTorch的GPU版本。可以通過以下命令安裝:

conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch 

或者使用pip

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

驗證安裝

安裝完成后,可以通過以下python代碼驗證PyTorch是否能夠使用GPU:

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

如果torch.cuda.is_available()返回True,則表示PyTorch已經成功配置并可以使用GPU。

請注意,具體的安裝步驟和版本可能會隨著軟件和驅動的更新而變化,建議在安裝前查閱最新的官方文檔或教程,以確保兼容性和穩定性。

? 版權聲明
THE END
喜歡就支持一下吧
點贊12 分享