Ubuntu中PyTorch與CUDA如何協同工作

在Ubuntu中，PyTorch與CUDA的協同工作主要依賴于正確安裝和配置兩者。以下是詳細的步驟和說明：

安裝顯卡驅動

首先，需要安裝NVIDIA的顯卡驅動。沒有安裝驅動的話要先安裝驅動，或者驅動版本太低的話要先進行升級?？梢允褂靡韵旅顏戆惭b推薦的顯卡驅動：

sudo ubuntu-drivers autoinstall

或者，如果你想安裝指定版本的驅動，可以使用：

sudo apt install nvidia-driver-555

安裝CUDA Toolkit

安裝完顯卡驅動后，接下來安裝CUDA Toolkit。這可以通過以下步驟完成：

1. 添加CUDA Toolkit的官方PPA（Personal Package Archive）：

sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt-get update

2. 安裝CUDA Toolkit：

sudo apt install cuda

3. 驗證CUDA安裝是否成功，可以使用以下命令：

nvcc -V

安裝cuDNN

最后，安裝cuDNN庫。這可以通過以下步驟完成：

1. 下載與CUDA版本兼容的cuDNN庫。例如，對于CUDA 11.8，可以從NVIDIA官網下載對應的cuDNN版本。

2. 解壓下載的cuDNN文件，并將解壓后的文件夾路徑添加到LD_LIBRARY_PATH環境變量中：

tar -xzvf cudnn-11.8.1-linux-x64-v8.6.0.70.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*
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

安裝PyTorch的GPU版本

最后，安裝PyTorch的GPU版本。這可以通過以下命令完成：

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

請將cu118替換為你安裝的CUDA版本。

驗證PyTorch與CUDA的協同工作

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

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

如果輸出為True，則表示PyTorch已經成功配置為使用CUDA。

以上步驟應該能夠幫助你在Ubuntu系統中成功配置PyTorch與CUDA的協同工作。如果在安裝過程中遇到任何問題，可以參考PyTorch和CUDA的官方文檔，或者在相關的技術社區尋求幫助。