圖像分類作為計算機視覺領(lǐng)域的核心任務(wù)，有著廣泛的應(yīng)用前景，如自動駕駛、醫(yī)療影像診斷、安防監(jiān)控等。PyTorch 憑借其強大的功能和靈活的操作，為開發(fā)者提供了一個高效構(gòu)建和訓(xùn)練圖像分類器的平臺。

一、數(shù)據(jù)準(zhǔn)備：構(gòu)建模型的基石

在訓(xùn)練分類器之前，我們需要準(zhǔn)備合適的訓(xùn)練數(shù)據(jù)。這里我們將使用經(jīng)典的 CIFAR10 數(shù)據(jù)集，它包含 10 個類別的彩色圖像，每個類別有 6000 張圖像，圖像大小為 32x32 像素。

使用 torchvision 加載 CIFAR10 數(shù)據(jù)集

import torch
import torchvision
import torchvision.transforms as transforms


## 數(shù)據(jù)預(yù)處理：將圖像轉(zhuǎn)換為張量，并進行標(biāo)準(zhǔn)化
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])


## 下載訓(xùn)練集和測試集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)


testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)


## 定義類別名稱
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

數(shù)據(jù)可視化

通過可視化部分訓(xùn)練數(shù)據(jù)，可以幫助我們更好地了解數(shù)據(jù)集的結(jié)構(gòu)和內(nèi)容。

import matplotlib.pyplot as plt
import numpy as np


## 定義一個函數(shù)用于顯示圖像
def imshow(img):
    img = img / 2 + 0.5  # 反標(biāo)準(zhǔn)化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


## 獲取一批訓(xùn)練數(shù)據(jù)
dataiter = iter(trainloader)
images, labels = next(dataiter)


## 顯示圖像
imshow(torchvision.utils.make_grid(images))


## 打印標(biāo)簽
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))

二、定義卷積神經(jīng)網(wǎng)絡(luò)：構(gòu)建分類器的核心

卷積神經(jīng)網(wǎng)絡(luò)（CNN）是處理圖像數(shù)據(jù)的主流網(wǎng)絡(luò)結(jié)構(gòu)，它通過卷積層自動提取圖像特征，能夠有效捕捉圖像中的空間信息。

定義 CNN 架構(gòu)

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 定義卷積層和池化層
        self.conv1 = nn.Conv2d(3, 6, 5)  # 輸入通道3，輸出通道6，卷積核大小5
        self.pool = nn.MaxPool2d(2, 2)   # 最大池化層，窗口大小2，步長2
        self.conv2 = nn.Conv2d(6, 16, 5) # 輸入通道6，輸出通道16，卷積核大小5
        # 定義全連接層
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)


    def forward(self, x):
        # 前向傳播過程
        x = self.pool(F.relu(self.conv1(x)))  # 卷積 + 激活 + 池化
        x = self.pool(F.relu(self.conv2(x)))  # 卷積 + 激活 + 池化
        x = x.view(-1, 16 * 5 * 5)            # 展平操作
        x = F.relu(self.fc1(x))               # 全連接 + 激活
        x = F.relu(self.fc2(x))               # 全連接 + 激活
        x = self.fc3(x)                       # 輸出層
        return x


net = Net()
print(net)

三、定義損失函數(shù)和優(yōu)化器：模型訓(xùn)練的指引

損失函數(shù)用于衡量模型預(yù)測結(jié)果與真實標(biāo)簽之間的差距，優(yōu)化器則負責(zé)根據(jù)損失函數(shù)的梯度信息更新模型參數(shù)。

import torch.optim as optim


## 使用交叉熵損失函數(shù)和隨機梯度下降優(yōu)化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

四、訓(xùn)練網(wǎng)絡(luò)：提升模型性能的關(guān)鍵

訓(xùn)練過程是模型學(xué)習(xí)數(shù)據(jù)特征、優(yōu)化參數(shù)的關(guān)鍵環(huán)節(jié)。我們需要多次迭代訓(xùn)練數(shù)據(jù)，逐步調(diào)整模型參數(shù)，以降低損失函數(shù)的值。

for epoch in range(2):  # 遍歷數(shù)據(jù)集多次
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 獲取輸入數(shù)據(jù)和標(biāo)簽
        inputs, labels = data


        # 清空梯度緩存
        optimizer.zero_grad()


        # 前向傳播 + 反向傳播 + 優(yōu)化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()


        # 打印統(tǒng)計信息
        running_loss += loss.item()
        if i % 2000 == 1999:    # 每 2000 個小批量打印一次
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0


print('Finished Training')

五、保存和加載模型：模型持久化與復(fù)用

訓(xùn)練完成后，我們可以將模型參數(shù)保存到文件中，以便后續(xù)加載和使用。

PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

加載模型參數(shù)：

net = Net()
net.load_state_dict(torch.load(PATH))

六、測試網(wǎng)絡(luò)：評估模型性能

在測試集上評估模型的性能，計算分類準(zhǔn)確率。

correct = 0
total = 0
with torch.no_grad():  # 在測試階段不需要計算梯度
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()


print(f'Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%')

七、在 GPU 上訓(xùn)練：加速模型訓(xùn)練

如果電腦配備 GPU，可以利用 GPU 加速模型訓(xùn)練過程，顯著提升訓(xùn)練速度。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)


## 將輸入數(shù)據(jù)和標(biāo)簽移動到 GPU 上
inputs, labels = inputs.to(device), labels.to(device)

八、總結(jié)

通過本教程，你已經(jīng)掌握了使用 PyTorch 訓(xùn)練圖像分類器的核心步驟，包括數(shù)據(jù)準(zhǔn)備、網(wǎng)絡(luò)定義、模型訓(xùn)練、性能評估以及 GPU 加速等關(guān)鍵技術(shù)。在編程獅平臺的進一步學(xué)習(xí)中，你可以嘗試以下方向：

• 探索更多數(shù)據(jù)集 ：除了 CIFAR10，還可以嘗試 ImageNet、MNIST 等其他知名數(shù)據(jù)集，挑戰(zhàn)不同難度的圖像分類任務(wù)。
• 優(yōu)化網(wǎng)絡(luò)結(jié)構(gòu) ：通過調(diào)整卷積層、池化層、全連接層的數(shù)量和參數(shù)，或者嘗試不同的網(wǎng)絡(luò)架構(gòu)（如 ResNet、AlexNet 等），提升模型性能。
• 學(xué)習(xí)高級技巧 ：深入了解數(shù)據(jù)增強、正則化、遷移學(xué)習(xí)等高級技巧，進一步提高模型的泛化能力和魯棒性。