CUDA 11.5 對應的 PyTorch 指南

CUDA 是 NVIDIA 公司推出的一個並行計算框架，它可為 GPU 提供良好的並行計算環境和編程接口，而 PyTorch 是一個基於 Python 的機器學習框架，這兩個框架的結合可使我們更高效地進行深度學習開發。本篇文章將從多個方面對 CUDA 11.5 對應的 PyTorch 做詳細的闡述。

一、 安裝和配置

安裝 CUDA 11.5 可以去官網下載安裝。對於 Python 環境的設置，可以通過 Anaconda 來管理。首先，新建並激活一個名為 pytorch 的虛擬環境，然後通過下面代碼安裝 PyTorch。

conda create -n pytorch python=3.8
conda activate pytorch
conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch

註：cudatoolkit 的版本需要和本機安裝的 CUDA 版本對應，這裡是 11.1 是因為 PyTorch 官方沒有在 11.5 上發布的相關版本，而在 11.1 上是可用的。

接着，在運行 PyTorch 代碼時，需要注意以下幾點：

1. 確認安裝的 PyTorch 版本是否支持當前 CUDA 的版本；

2. 確認 GPU 是否支持 CUDA；

3. 通過下面的代碼檢查 PyTorch 是否安裝成功。

import torch

print(torch.cuda.is_available())  # 檢查 GPU 是否支持 CUDA
print(torch.__version__)          # 檢查 PyTorch 版本號

二、 數據預處理

在進行深度學習訓練時，數據預處理是一個非常重要的環節。PyTorch 提供了多種方式來進行數據預處理。對於圖像數據，我們可以使用 torchvision 來預處理數據，並將其轉換為 PyTorch 中的 Tensor 數據類型。

下面是一個示例代碼，它使用了 torchvision.transforms 中的 Compose、Resize、ToTensor 函數，將圖像數據轉換為 PyTorch 中的 Tensor 數據類型。

from torchvision.transforms import Compose, Resize, ToTensor
from PIL import Image

dataset_dir = 'your dataset directory'
dataset = []

# 加載圖像並進行預處理
transform = Compose([Resize((224, 224)), ToTensor()])

for img_path in os.listdir(dataset_dir):
    img = Image.open(os.path.join(dataset_dir, img_path)).convert('RGB')
    x = transform(img)
    dataset.append(x)

上述代碼中，Componse 函數用於將多個預處理操作組合在一起，Resize 函數用於調整圖像大小，ToTensor 函數用於將圖像數據轉換為 Tensor 數據類型。在循環中，我們使用 PIL 庫來加載圖像數據，將其轉換為 PyTorch 中的 Tensor 數據類型，並將其添加到 dataset 列表中。

三、 神經網絡搭建

在 PyTorch 中，我們可以使用 nn.Module 類來構建神經網絡模型。該類提供了一些有用的函數，例如 forward 函數，用於定義輸入和輸出之間的計算。

下面是一個示例代碼，它演示如何使用 nn.Module 類構建一個簡單的神經網絡，並使用該網絡對圖像進行分類。

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(32 * 56 * 56, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 32 * 56 * 56)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

net = Net()

# 在 GPU 上運行模型
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
net.to(device)

上述代碼中，我們首先定義了一個 Net 類，並在 __init__ 構造函數中定義了神經網絡的結構，包括卷積層、池化層和全連接層。在 forward 函數中，我們定義了輸入和輸出之間的計算。接下來，我們創建了一個 net 實例，並通過 to 函數將其移至 GPU 上進行訓練。

四、 模型訓練

在神經網絡模型搭建完成後，我們需要對其進行訓練，並評估模型的性能。在 PyTorch 中，我們可以使用 nn.CrossEntropyLoss 函數來計算分類損失，並使用反向傳播算法進行模型參數的更新。

下面是一個示例代碼，它展示了如何使用 PyTorch 進行模型訓練。

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(2):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 100))

上述代碼中，我們首先定義了損失函數和優化器。在 for 循環中，我們遍歷訓練數據集，將輸入和標籤移至 GPU 上，並使用 optimizer.zero_grad() 函數清空梯度信息；然後我們計算輸出，根據損失函數計算損失並進行反向傳播，最後使用 optimizer.step() 函數更新模型參數。運行過程中，我們統計並輸出損失值。

五、 模型評估

在模型訓練後，我們需要對其進行評估，以了解模型的性能如何。在 PyTorch 中，我們可以使用模型的 eval() 函數來將其設置為評估模式，並使用 nn.functional.softmax 函數計算對數概率值。

下面是一個示例代碼，它演示了如何使用 PyTorch 進行模型評估。

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))

上述代碼中，我們使用 eval() 函數將模型設置為評估模式，並遍歷測試數據集，使用 nn.functional.softmax 函數計算對數概率值，然後使用 torch.max 函數計算出最大值及其對應的位置，最後計算準確率。

總結

本篇文章從安裝和配置、數據預處理、神經網絡搭建、模型訓練、模型評估等多個方面詳細介紹了 CUDA 11.5 對應的 PyTorch。希望本篇文章可以對深度學習開發者有所幫助。