CNN手寫數字識別demo——pyTorch代碼實現

import time
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt

EPOCH = 1                           # 迭代次數
BATCH_SIZE = 50                     # 每一批訓練的數據大小
LR = 0.001                          # 學習率
DOWNLOAD_MNIST = False              # 是否需要下載數據集，首次運行需要把該參數指定爲True

'''
torchvisions 是一個數據集的庫，下載torch的時候一起下載的
MNIST 是手寫數字的數據集
root 是數據集的根目錄
train爲True 代表的是下載訓練數據，爲False代表下載的是測試數據
transform 是指定數據格式，ToTensor是將圖片轉化爲0到1之間的Tensor類型的數據
download 表示是否需要下載數據集，如果已經下載過，就設置爲False
'''
train_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=DOWNLOAD_MNIST
)

# # 查看數據size
# print(train_data.data.size())
# # 查看標籤size
# print(train_data.targets.size())
# # 在畫布上展示圖片
# plt.imshow(train_data.data[0].numpy(), cmap='gray')
# # 設置標題，%i 代表轉爲有符號的十進制
# plt.title('%i'%train_data.targets[0])
# plt.show()

#  訓練集分批
train_loader = Data.DataLoader(
    dataset=train_data,
    batch_size=BATCH_SIZE,
    shuffle=True,
)

# 下載測試集數據
test_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=False
)

'''
這裏的test_data.data的size是(10000, 28, 28)
維度一 表示測試集總數
維度二 表示第多少個像素的橫座標對應的RGB值（這裏已從三層(xxx,xxx,xxx)轉爲一層(k), k的值是0到1之間）
維度三 與參數二相似，表示的是縱座標
添加維度(圖片特徵的高度)並切片後，test_x的size是(2000, 1, 28, 28)
PyTorch圖層以Torch.FloatTensor作爲輸入，所以這裏要指定一下torch.float32，否則在後續的訓練中會報錯。
最後除255是因爲像素點是0-255範圍的，這樣可以做歸一化處理，把特徵值固定在0-1範圍區間。如果不做
歸一化處理，會在某些情況下（比如RNN神經網絡）出現欠擬合的問題。
'''
test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[:2000]/255
# 獲取正確分類的結果
test_y = test_data.targets[:2000]


class CNN(nn.Module):

    def __init__(self):
        super(CNN, self).__init__()
        '''
        對於nn.Conv2d中的padding參數，單獨解釋一下。
        把每個像素點當成一個特徵，圖像爲28*28個像素，以長寬爲5，步長爲1的過濾器去掃描時，x軸方向，
        第一次掃描的是1~5個像素點（作爲下一層神經網絡的第1個特徵），第二次掃描的是第2~6個像素點（作爲下一層
        神經網絡的第2個特徵）……，最後一次掃描的是第24~28個像素點（作爲下一層神經網絡的第24個特徵），會發現下一層
        比上一層少了4個特徵。y軸方向也是同理。
        padding=2 表示在圖像的四周都分別填充2個像素點。這樣圖像就變成了32*32個像素。
        此時按上述方式進行，會發現下一層的特徵數量爲28，這樣就能保證兩個神經網絡層之間的特徵數量相同。
        所以，想要 con2d 出來的圖片長寬沒有變化, padding=(kernel_size-stride)/2
        '''
        '''
        關於nn.MaxPool2d(kernel_size=2)
        把卷積後並經過激活函數的數據進行池化，MaxPooLing的方式池化（選取最大的一個像素點），
        池化的作用是壓縮圖片，篩選重要的特徵信息。
        可以把它看成一個過濾器，kernel_size是過濾器的大小
        '''
        # 卷積層 主要是用做過濾
        self.conv1 = nn.Sequential(
            nn.Conv2d(            # 過濾器
                in_channels=1,    # 輸入的高度，因爲已經從RGB三層轉爲一層，所以這裏是1
                out_channels=16,  # 過濾器的個數，也是輸出的高度，每個過濾器收集一塊圖像區域的信息，作爲一層圖像高度信息
                kernel_size=5,    # 過濾器的大小(長寬)
                stride=1,         # 過濾器掃描時的步長
                padding=2         # 如果想要 con2d 出來的圖片長寬沒有變化, padding=(kernel_size-stride)/2
            ),
            nn.ReLU(),            # 激活函數
            nn.MaxPool2d(kernel_size=2)  # 池化，可以把它看成一個過濾器，kernel_size是過濾器的大小
        )
        self.conv2 = nn.Sequential(
            # 第一次卷積後，輸出的高度爲16，所以這裏輸入的高度是16
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        '''
        把經過兩次卷積和池化的信息輸入給全連接神經層
        32 * 7 * 7分別代表圖片的高度和長、寬，是輸入特徵，10是輸出特徵（全類別的概率分佈）。
        原本的圖片是28個像素，經過一次卷積，但是因爲定義了padding，卷積後圖片長寬沒有改變。在經過池化時，由於
        kernel_size=2，也就是說從2*2的像素塊中提取一個最大像素的特徵，既單獨觀察x軸，是從兩個像素中抽取了一個。
        單獨觀察y軸也是同理。所以經過第一次池化時，像素點從28*28變成了14*14。再經過第二次卷積和激活函數，
        像素點不變，經過第二次池化，像素點變成了7*7。
        '''
        self.out = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)           # torch.Size([50, 32, 7, 7])
        x = x.view(x.size(0), -1)   # torch.Size([50, 1568])，保留batch、把高、長、寬拉平，變成一維數據，相當於每次批訓練 50 張圖片，每張圖片1568個特徵
        output = self.out(x)
        return output

# 創建cnn模型對象
cnn = CNN()
# 指定分類器
optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)
# 指定誤差計算的方式，一般來說，迴歸問題使用均方差，分類問題使用交叉熵
loss_func = nn.CrossEntropyLoss()

start_time = time.time()

for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):
        b_x = Variable(x)
        b_y = Variable(y)

        output = cnn(x)                                             # 傳入測試數據，得出預測值
        loss = loss_func(output, b_y)                               # 對比預測值與真實值，得出loss值
        optimizer.zero_grad()                                       # 優化器梯度歸零
        loss.backward()                                             # 進行誤差反向傳遞
        optimizer.step()                                            # 將參數更新值施加到 cnn 的 parameters 上

        if step % 50 == 0:
            test_output = cnn(test_x)                               # torch.Size([2000, 10])
            pred_y = torch.max(test_output, 1)[1]                   # torch.Size([2000])
            accuracy = (sum(pred_y == test_y)).numpy() / test_y.size(0)     # 算出識別精度
            print('Epoch：', epoch, '| train loss：%.4f' % loss.data.numpy(), '| test accuracy：%.2f' % accuracy)

# 使用訓練好的模型，預測前十個數據，然後和真實值對比
test_output = cnn(test_x[:10])
pred_y = torch.max(test_output, 1)[1]
print(pred_y, 'prediction number')
print(test_y[:10], 'real number')

print('耗時：', time.time() - start_time)

輸出結果如下：

Epoch： 0 | train loss：2.3105 | test accuracy：0.26
Epoch： 0 | train loss：0.4869 | test accuracy：0.81
Epoch： 0 | train loss：0.4686 | test accuracy：0.87
Epoch： 0 | train loss：0.1397 | test accuracy：0.90
Epoch： 0 | train loss：0.0822 | test accuracy：0.93
Epoch： 0 | train loss：0.0621 | test accuracy：0.94
Epoch： 0 | train loss：0.1441 | test accuracy：0.95
Epoch： 0 | train loss：0.2057 | test accuracy：0.95
Epoch： 0 | train loss：0.1435 | test accuracy：0.95
Epoch： 0 | train loss：0.1141 | test accuracy：0.96
Epoch： 0 | train loss：0.0555 | test accuracy：0.96
Epoch： 0 | train loss：0.1598 | test accuracy：0.95
Epoch： 0 | train loss：0.1003 | test accuracy：0.97
Epoch： 0 | train loss：0.0447 | test accuracy：0.97
Epoch： 0 | train loss：0.0247 | test accuracy：0.96
Epoch： 0 | train loss：0.0342 | test accuracy：0.97
Epoch： 0 | train loss：0.0354 | test accuracy：0.97
Epoch： 0 | train loss：0.0867 | test accuracy：0.97
Epoch： 0 | train loss：0.0299 | test accuracy：0.97
Epoch： 0 | train loss：0.0778 | test accuracy：0.97
Epoch： 0 | train loss：0.0561 | test accuracy：0.97
Epoch： 0 | train loss：0.0235 | test accuracy：0.97
Epoch： 0 | train loss：0.1320 | test accuracy：0.98
Epoch： 0 | train loss：0.0668 | test accuracy：0.97
tensor([7, 2, 1, 0, 4, 1, 4, 9, 5, 9]) prediction number
tensor([7, 2, 1, 0, 4, 1, 4, 9, 5, 9]) real number
耗時： 89.75633668899536

可以看到，loss值逐漸降低，精度逐漸提升，最後穩定在97%左右的準確率。而且最後的預測值與真實值相同，模型還是不錯的。但是由於是使用CPU訓練（本機是i5-8300H），耗時較長，跑了差不多一分半鐘的時間。