在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版本
最后,使用conda或pip安裝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。
請注意,具體的安裝步驟和版本可能會隨著軟件和驅動的更新而變化,建議在安裝前查閱最新的官方文檔或教程,以確保兼容性和穩定性。