如何在Ubuntu上安裝PyTorch GPU版本

在Ubuntu上安裝PyTorch GPU版本需要以下幾個步驟：

安裝NVIDIA驅動：確保你的系統已經安裝了與你的GPU兼容的NVIDIA驅動。你可以通過以下命令來檢查是否已經安裝了NVIDIA驅動：
```
nvidia-smi
```
如果顯示了GPU信息，說明驅動已經安裝。如果沒有，你需要先安裝驅動。
添加NVIDIA CUDA Toolkit： PyTorch GPU版本依賴于CUDA Toolkit。你可以從NVIDIA官網下載并安裝適合你系統的CUDA Toolkit。以下是安裝CUDA Toolkit的步驟：
- 訪問NVIDIA CUDA Toolkit下載頁面。
- 選擇適合你Ubuntu版本的CUDA Toolkit進行下載。
- 下載完成后，按照安裝向導進行安裝。例如，對于CUDA 11.7，你可以使用以下命令：
```
sudo dpkg -i cuda-repo-ubuntu2004-11-7-local_11.7.0-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2004-11-7-local/7fa2af80.pub
sudo apt-get update
sudo apt-get install cuda
```
安裝cuDNN： cuDNN是用于深度神經網絡的GPU加速庫。你需要從NVIDIA官網下載并安裝與你的CUDA版本兼容的cuDNN庫。以下是安裝cuDNN的步驟：
- 訪問NVIDIA cuDNN下載頁面。
- 注冊并登錄NVIDIA開發者賬號。
- 下載適合你CUDA版本的cuDNN庫。
- 解壓下載的文件并將文件復制到CUDA目錄中。例如：
```
tar -xzvf cudnn-11.7-linux-x64-v8.4.1.50.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*
```
設置環境變量：為了確保PyTorch能夠找到CUDA和cuDNN庫，你需要設置一些環境變量。你可以在~/.bashrc文件中添加以下內容：
```
export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
```
然后運行以下命令使環境變量生效：
```
source ~/.bashrc
```
安裝PyTorch GPU版本：你可以使用pip來安裝PyTorch GPU版本。首先，確保你已經安裝了pip。然后，運行以下命令來安裝PyTorch：
```
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
```
這里的cu117表示CUDA 11.7版本。如果你使用的是其他版本的CUDA，請相應地更改URL。

完成以上步驟后，你應該能夠在Ubuntu上成功安裝并運行PyTorch GPU版本。你可以通過以下命令來驗證安裝是否成功：

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

如果輸出為True，說明PyTorch GPU版本已經成功安裝并可以使用。

最新問答

相關標簽