Ubuntu上PyTorch與CUDA如何協同工作

在Ubuntu系統上，PyTorch與CUDA的協同工作主要依賴于NVIDIA提供的CUDA工具包和驅動程序，以及PyTorch的CUDA支持庫。以下是它們協同工作的基本步驟：

1. 安裝NVIDIA驅動程序

首先，確保你的系統已經安裝了兼容的NVIDIA顯卡驅動程序。你可以從NVIDIA官方網站下載并安裝最新的驅動程序。

sudo apt update
sudo apt install nvidia-driver-<version>

2. 安裝CUDA工具包

接下來，安裝CUDA工具包。你可以從NVIDIA官方網站下載適用于Ubuntu的CUDA Toolkit安裝包，并按照官方指南進行安裝。

# 下載CUDA Toolkit安裝包
wget https://developer.download.nvidia.com/compute/cuda/<version>/cuda_<version>_linux.run

# 安裝CUDA Toolkit
sudo sh cuda_<version>_linux.run

安裝過程中，確保選擇安裝CUDA Toolkit和驅動程序。

3. 安裝cuDNN庫

cuDNN（CUDA Deep Neural Network library）是一個用于深度神經網絡的GPU加速庫。你需要從NVIDIA官方網站下載并安裝cuDNN庫。

# 下載cuDNN庫
wget https://developer.nvidia.com/compute/machine-learning/cudnn/secure/<version>/prod/<version>_linux_x64_v<version>.tgz

# 解壓并安裝cuDNN庫
tar -xzvf <version>_linux_x64_v<version>.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*

4. 安裝PyTorch與CUDA支持

使用pip安裝PyTorch，并確保安裝的版本支持CUDA。

# 安裝PyTorch與CUDA支持
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu<version>

將<version>替換為你安裝的CUDA版本號。例如，如果你安裝的是CUDA 11.7，命令應為：

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

5. 驗證安裝

驗證PyTorch是否能夠檢測到CUDA設備。

import torch

# 檢查是否有可用的GPU
print(torch.cuda.is_available())

# 獲取GPU設備數量
print(torch.cuda.device_count())

# 獲取當前GPU設備名稱
print(torch.cuda.get_device_name(0))

如果輸出顯示True，并且能夠列出GPU設備名稱，則說明PyTorch與CUDA已經成功協同工作。

6. 配置環境變量

確保CUDA和cuDNN的路徑已經添加到系統的環境變量中。編輯~/.bashrc文件，添加以下內容：

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

然后運行以下命令使更改生效：

source ~/.bashrc

通過以上步驟，你可以在Ubuntu系統上成功安裝并配置PyTorch與CUDA，使其能夠協同工作。