前言：此前分享了卷積神經網絡相關知識，今天實戰下：搭建一個卷積神經網絡來實現圖像分類任務。

一、了解CIFAR-10數據集

CIFAR-10數據集5萬張訓練圖像、1萬張測試圖像、10個類別、每個類別有6k個圖像，圖像大小32×32×3。下圖列舉了10個類，每一類隨機展示了10張圖片：

PyTorch 中的 torchvision.datasets 計算機視覺模塊封裝了 CIFAR10 數據集，如果需要使用可以直接導入。

導入代碼：

from torchvision.datasets import CIFAR10

二、案例之導包

import torch
import torch.nn as nn
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor  # pip install torchvision -i https://mirrors.aliyun.com/pypi/simple/
import torch.optim as optim
from torch.utils.data import DataLoader
import time
import matplotlib.pyplot as plt
from torchsummary import summary# 每批次樣本數
BATCH_SIZE = 8

三、案例之創建數據集

# 1. 數據集基本信息
def create_dataset():# 加載數據集:訓練集數據和測試數據# ToTensor: 將image（一個PIL.Image對象）轉換為一個Tensortrain = CIFAR10(root='data', train=True, transform=ToTensor())valid = CIFAR10(root='data', train=False, transform=ToTensor())# 返回數據集結果return train, validif __name__ == '__main__':# 數據集加載train_dataset, valid_dataset = create_dataset()# 數據集類別print("數據集類別:", train_dataset.class_to_idx)# 數據集中的圖像數據print("訓練集數據集:", train_dataset.data.shape)print("測試集數據集:", valid_dataset.data.shape)# 圖像展示plt.figure(figsize=(2, 2))plt.imshow(train_dataset.data[1])plt.title(train_dataset.targets[1])plt.show()

運行結果：

數據集類別: {'airplane': 0, 'automobile': 1, 'bird': 2, 'cat': 3, 'deer': 4, 'dog': 5, 'frog': 6, 'horse': 7, 'ship': 8, 'truck': 9}
訓練集數據集: (50000, 32, 32, 3)
測試集數據集: (10000, 32, 32, 3)

圖像：

四、案例之搭建神經網絡（模型構建）

需要搭建的CNN網絡結構如下:

我們要搭建的網絡結構如下:

1. 輸入形狀: 32x32；
2. 第一個卷積層輸入 3 個 Channel, 輸出 6 個 Channel, Kernel Size 為: 3x3；
3. 第一個池化層輸入 30x30, 輸出 15x15, Kernel Size 為: 2x2, Stride 為: 2；
4. 第二個卷積層輸入 6 個 Channel, 輸出 16 個 Channel, Kernel Size 為 3x3；
5. 第二個池化層輸入 13x13, 輸出 6x6, Kernel Size 為: 2x2, Stride 為: 2；
6. 第一個全連接層輸入 576 維, 輸出 120 維；
7. 第二個全連接層輸入 120 維, 輸出 84 維；
8. 最后的輸出層輸入 84 維, 輸出 10 維。

我們在每個卷積計算之后應用 relu 激活函數來給網絡增加非線性因素。

# 模型構建
class ImageClassification(nn.Module):# 定義網絡結構def __init__(self):super(ImageClassification, self).__init__()# 定義網絡層：卷積層+池化層# 第一個卷積層, 輸入圖像為3通道,輸出特征圖為6通道,卷積核3*3self.conv1 = nn.Conv2d(3, 6, stride=1, kernel_size=3)# 第一個池化層, 核寬高2*2self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)# 第二個卷積層, 輸入圖像為6通道,輸出特征圖為16通道,卷積核3*3self.conv2 = nn.Conv2d(6, 16, stride=1, kernel_size=3)# 第二個池化層, 核寬高2*2self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)# 全連接層# 第一個隱藏層 輸入特征576個(一張圖像為16*6*6), 輸出特征120個self.linear1 = nn.Linear(576, 120)# 第二個隱藏層self.linear2 = nn.Linear(120, 84)# 輸出層self.out = nn.Linear(84, 10)# 定義前向傳播def forward(self, x):# 卷積+relu+池化x = torch.relu(self.conv1(x))x = self.pool1(x)# 卷積+relu+池化x = torch.relu(self.conv2(x))x = self.pool2(x)# 將特征圖做成以為向量的形式：相當于特征向量 全連接層只能接收二維數據集# 由于最后一個批次可能不夠8，所以需要根據批次數量來改變形狀# x[8, 16, 6, 6] --> [8, 576] -->8個樣本,576個特征# x.size(0): 第1個值是樣本數 行數# -1：第2個值由原始x剩余3個維度值相乘計算得到 列數(特征個數)x = x.reshape(x.size(0), -1)# 全連接層x = torch.relu(self.linear1(x))x = torch.relu(self.linear2(x))# 返回輸出結果return self.out(x)if __name__ == '__main__':# 模型實例化model = ImageClassification()summary(model, input_size=(3,32,32), batch_size=1)

運行結果：

五、案例之編寫訓練函數（訓練模型）

在訓練時，使用多分類交叉熵損失函數，Adam 優化器。具體實現代碼如下:

def train(model, train_dataset):# 構建數據加載器dataloader = DataLoader(train_dataset, batch_size=10, shuffle=True)criterion = nn.CrossEntropyLoss() # 構建損失函數optimizer = optim.Adam(model.parameters(), lr=1e-3) # 構建優化方法epoch = 100  # 訓練輪數for epoch_idx in range(epoch):sum_num = 0   # 樣本數量total_loss = 0.0  # 損失總和correct = 0  # 預測正確樣本數start = time.time()  # 開始時間# 遍歷數據進行網絡訓練for x, y in dataloader:model.train()output = model(x)loss = criterion(output, y)  # 計算損失optimizer.zero_grad()  # 梯度清零loss.backward()  # 反向傳播optimizer.step()  # 參數更新correct += (torch.argmax(output, dim=-1) == y).sum()  # 計算預測正確樣本數# 計算每次訓練模型的總損失值 loss是每批樣本平均損失值total_loss += loss.item()*len(y)  # 統計損失和sum_num += len(y)print('epoch:%2s loss:%.5f acc:%.2f time:%.2fs' %(epoch_idx + 1,total_loss / sum_num,correct / sum_num,time.time() - start))# 模型保存torch.save(model.state_dict(), 'model/image_classification.pth')#聯合上面代碼一起運行本代碼
if __name__ == '__main__':# 數據集加載train_dataset, valid_dataset = create_dataset()# 模型實例化model = ImageClassification()# 模型訓練train(model,train_dataset)

運行結果：

epoch: 1 loss:1.67102 acc:0.38 time:26.23s
epoch: 2 loss:1.35650 acc:0.51 time:27.63s
epoch: 3 loss:1.22355 acc:0.57 time:31.10s
epoch: 4 loss:1.14639 acc:0.59 time:66.37s
epoch: 5 loss:1.09468 acc:0.61 time:40.38s
。。。。。。

六、案例之編寫預測函數（模型測試）

當已經訓練好模型（model），并保存了模型參數（model.state_dict()），可直接實例化模型，并加載訓練好的模型參數，然后對測試集中的1萬條樣本進行預測，查看模型在測試集上的準確率。

def eval(valid_dataset):# 構建數據加載器dataloader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False)# 加載模型并加載訓練好的權重model = ImageClassification()model.load_state_dict(torch.load('model/image_classification.pth'))# 模型切換評估模式, 如果網絡模型中有dropout/BN等層, 評估階段不進行相應操作model.eval()# 計算精度total_correct = 0total_samples = 0# 遍歷每個batch的數據，獲取預測結果，計算精度for x, y in dataloader:output = model(x)total_correct += (torch.argmax(output, dim=-1) == y).sum()total_samples += len(y)# 打印精度print('Acc: %.2f' % (total_correct / total_samples))if __name__ == '__main__':train_dataset, valid_dataset = create_dataset()eval(valid_dataset)

運行結果：

Acc: 0.57

最后，大家還可以通過調整lr（學習率）、神經元失活（dropout）、增加神經網絡層數等方式來調整模型，提升acc，各看本領吧！

今天的分享到此結束。