【PyTorch】MNIST手寫數字識別實踐

一、數據集

1.1 數據集介紹

MNIST數據集（官網）被分成兩部分：60000行的訓練數據集（mnist.train）和10000行的測試數據集（mnist.test）。
每一個MNIST數據單元由兩部分組成：一張包含手寫數字的圖片和一個對應的標籤。每一張圖片包含28X28個像素點。我們可以用一個數字數組來表示這張圖片：

1.2 數據集下載和加載

1.2.1 在線加載

在Pycharm中新建項目並命名爲pytorch_mnist，在根目錄下新建dataset.py文件用於存放數據集相關代碼，代碼說明請看註釋。完整代碼如下：

import torch
from torchvision import datasets, transforms

class Dataset():
    def __init__(self):
        # 注意這是python 2.0的寫法
        super(Dataset, self).__init__()
        # python 3.0+可省略super()中的參數
        #  super().__init__()

        # 一個批次加載的圖片數量
        self.batch_size = 64
        # 數據預處理
        # Compose用於將多個transfrom組合起來
        # ToTensor()將像素轉換爲tensor，並做Min-max歸一化，即x'=x-min/max-min
        # 相當於將像素從[0,255]轉換爲[0,1]
        # Normalize()用均值和標準差對圖像標準化處理 x'=(x-mean)/std，加速收斂的作用
        # 這裏0.131是圖片的均值，0.308是方差，通過對原始圖片進行計算得出
        # 想偷懶的話可以直接填Normalize([0.5], [0.5])
        # 另外多說一點，因爲MNIST數據集圖片是灰度圖，只有一個通道，因此這裏的均值和方差都只有一個值
        # 若是普通的彩色圖像，則應該是三個值，比如Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        self.transforms = transforms.Compose([transforms.ToTensor(),
                                              transforms.Normalize([0.131], [0.308])])
        # 下載數據集
        # 訓練數據集 train=True
        # './data/mnist'是數據集存放的路徑，可自行調整
        # download=True表示叫pytorch幫我們自動下載
        self.data_train = datasets.MNIST('./data/mnist',
                                    train=True,
                                    transform=self.transforms,
                                    download=True
                                    )

        # 測試數據集 train=False
        self.data_test = datasets.MNIST('./data/mnist',
                                   train=False,
                                   transform=self.transforms,
                                   download=True
                                   )

        # 加載數據集
        # shuffle=True表示加載時打亂圖片順序，有一定的防止過擬合效果
        self.loader_train = torch.utils.data.DataLoader(self.data_train,
                                                   batch_size=self.batch_size,
                                                   shuffle=True)
        # 測試集就不需要打亂了，因此shuffle=False
        self.loader_test = torch.utils.data.DataLoader(self.data_test,
                                                  batch_size=self.batch_size,
                                                  shuffle=False)

這裏我們可以嘗試對數據集進行加載，在dataset.py繼續添加如下代碼：

if __name__ == '__main__':
    dst = Dataset()

這段代碼是對Dataset類進行了實例化，會自動調用Dataset類中的構造方法__init__，並執行數據下載和加載代碼。運行效果如下：

數據下載完成後，觀察我們的項目目錄會發現數據集已經下載到本地了，如下：

看到這裏說明你的數據加載已經沒問題了，走向成功的第一步已經ok。

1.2.2 本地加載

如果你已經將MNIST下載到了本地，還是用上面的datasets.MNIST來加載數據的話，則需將下載好的數據集按如下目錄存放，不然會出問題。

其中data/mnist是在datasets.MNIST填寫的下載路徑，需要修改dataset.py中的部分代碼，將download=False即可，如下:

datasets.MNIST('./data/mnist',
               train=True,
               transform=self.transforms,
               download=False)

好了本地數據集加載已經搞定啦。

1.3 數據預覽

前面我們已經下載好了數據集並寫好了數據加載器，現在實現數據預覽，在項目根目錄新建main.py文件，添加如下代碼：

import torch
import torchvision
from dataset import Dataset
import cv2

dst = Dataset()
def show():
   imgs, labels = next(iter(dst.loader_train))
   # 將一個批次的圖拼成雪碧圖展示
   # 此時img的維度爲[channel, height, width]
   img = torchvision.utils.make_grid(imgs)
   # 轉換爲numpy數組並調整維度爲[height, width, channel]
   # 因爲下面的cv2.imshow()方法接受的數據的維度應該這樣
   img = img.numpy().transpose(1, 2, 0)
   # 因爲之前預處理對數據做了標準差處理
   # 這裏需要逆過程來恢復
   img = img * 0.308 + 0.131
   # 打印圖片對應標籤
   print(labels)
   # 展示圖片
   cv2.imshow('mnist', img)
   # 等待圖片關閉
   key_pressed = cv2.waitKey(0)

if __name__ == '__main__':
   show()

需要注意這裏我們使用了opencv的cv2庫進行數據預覽，需要預先下載該模塊，具體下載方法這裏就不細說了。
運行結果如下：

看到這裏，說明我們的數據加載也沒問題了，接下載可以開始構建我們的網絡模型了。

1.4 網絡模型搭建

這裏我們使用兩種網絡模型，第一種由我們自己定義，第二種借鑑經典網絡LeNet5。
在項目根目錄下新建文件夾models，並在該目錄下新建MyCNN.py文件。這裏我們構建一個只有一個卷積層，一個池化層，和三個全連接層的簡陋網絡。

import torch.nn as nn

class MyCNN(nn.Module):
   # 因爲分10類，設置n_classes=10
   def __init__(self, n_classes=10):
       super(MyCNN, self).__init__()

       # 關於pytorch中網絡層次的定義有幾種方式，這裏用的其中一種，用nn.Sequential()進行組合
       # 另外還可用有序字典OrderedDict進行定義
       # 再或者不使用nn.Sequential()進行組合，而是每一層單獨定義
       # 看具體需求和個人愛好
       self.features = nn.Sequential(
           # 輸入28×28，灰度圖in_channels=1
           nn.Conv2d(in_channels=1, out_channels=16, kernel_size=7, stride=1, padding=0),  # 輸出22×22
           nn.ReLU(inplace=True),
           nn.MaxPool2d(kernel_size=2, stride=2),  # 輸出11×11
       )

       self.classifier = nn.Sequential(
           nn.Linear(16 * 11 * 11, 160),
           nn.Linear(160, 80),
           nn.Linear(80, n_classes)
       )

   # 定義前向傳播函數
   def forward(self, x):
       # 將輸入送入卷積核池化層
       out = self.features(x)
       print(out.shape)
       # 這裏需要將out扁平化，展開成一維向量
       # 具體可慘開view()的用法
       out = out.view(out.size()[0], -1)
       # 將卷積和池化後的結果送入全連接層
       out = self.classifier(out)

       return out

if __name__ == '__main':
   print(MyCNN())

在models目錄下 新建LeNet5.py文件，添加如下代碼：

import torch.nn as nn
from collections import OrderedDict

class LeNet5(nn.Module):
   def __init__(self):
       super(LeNet5, self).__init__()

       # 這裏演示用OrderedDict定義網絡模型結果
       # 注意每一層的命名不要重複，不然重複的會不起作用
       self.conv = nn.Sequential(OrderedDict([
           ('c1', nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2)),
           ('relu1', nn.ReLU()),
           ('s2', nn.MaxPool2d(kernel_size=2, stride=2)),
           ('c3', nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)),
           ('relu3', nn.ReLU()),
           ('s4', nn.MaxPool2d(2, 2))
       ]))
       self.fc = nn.Sequential(OrderedDict([
           ('f6', nn.Linear(16 * 5 * 5, 120)),
           ('relu6', nn.ReLU()),
           ('f7', nn.Linear(120, 84)),
           ('relu7', nn.ReLU()),
           ('f8', nn.Linear(84, 10)),
       ]))

   def forward(self, x):
       x = self.conv(x)
       x = x.view(x.size()[0], -1)
       x = self.fc(x)
       return x

至此，我們的自定義網絡模型已經搭建好了，接下來開始訓練。

1.5 模型訓練

在main.py中添加如下代碼：請忽略之前已經寫過的代碼show()

import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
from dataset import Dataset
from models.MyCNN import MyCNN
from models.LeNet5 import LeNet5
import cv2

dst = Dataset()
# 模型實例化
my_model = MyCNN()
# my_model = LeNet5()

# 定義損失函數
loss_fn = nn.CrossEntropyLoss()
# 定義優化器
optimizer = optim.Adam(my_model.parameters(), lr=1e-4)

# 將模型的所有參數拷貝到到GPU上
if torch.cuda.is_available():
   my_model = my_model.cuda()

# 爲了節省時間成本，這裏我們只訓練5個epoch
# 可以根據實際情況進行調整
def train(epoches=5):
   for epoch in range(1, epoches + 1):
       print('Epoch {}/{}'.format(epoch, epoches))
       print('-' * 20)

       # 損失值
       running_loss = 0.0
       # 預測的正確數
       running_correct = 0

       for batch, (imgs, labels) in enumerate(dst.loader_train, 1):
           if torch.cuda.is_available():
               # 獲取輸入數據X和標籤Y並拷貝到GPU上
               # 注意有許多教程再這裏使用Variable類來包裹數據以達到自動求梯度的目的，如下
               # Variable(imgs)
               # 但是再pytorch4.0之後已經不推薦使用Variable類，Variable和tensor融合到了一起
               # 因此我們這裏不需要用Variable
               # 若我們的某個tensor變量需要求梯度，可以用將其屬性requires_grad=True,默認值爲False
               # 如，若X和y需要求梯度可設置X.requires_grad=True，y.requires_grad=True
               # 但這裏我們的X和y不需要進行更新，因此也不用求梯度
               X, y = imgs.cuda(), labels.cuda()
           else:
               X, y = imgs, labels

           # 將輸入X送入模型進行訓練
           outputs = my_model(X)
           # torch.max()返回兩個字，其一是最大值，其二是最大值對應的索引值
           # 這裏我們用y_pred接收索引值
           _, y_pred = torch.max(outputs.detach(), dim=1)
           # 在求梯度前將之前累計的梯度清零，以免影響結果
           optimizer.zero_grad()
           # 計算損失值
           loss = loss_fn(outputs, y)
           # 反向傳播
           loss.backward()
           # 參數更新
           optimizer.step()

           # 計算一個批次的損失值和
           running_loss += loss.detach().item()
           # 計算一個批次的預測正確數
           running_correct += torch.sum(y_pred == y)

           # 打印訓練結果
           if batch == len(dst.loader_train):
               print(
                   'Batch {batch}/{iter_times},Train Loss:{loss:.2f},Train Acc:{correct}/{lens}={acc:.2f}%'.format(
                       batch=batch,
                       iter_times=len(dst.loader_train),
                       loss=running_loss / batch,
                       correct=running_correct.item(),
                       lens=32 * batch,
                       acc=100 * running_correct.item() / (dst.batch_size * batch)
                   ))
               print('-' * 20)

       # 保存我們訓練好的模型
       if epoch == epoches:
           torch.save(my_model, 'models/MyModels.pth')
           print('Saving models/MyModels.pth')
if __name__ == '__main__':
   train()

查看訓練結果：

LeNet5的訓練結果：

沒想到我們自定義的網絡模型能訓練的準確率比LeNet5高吶（手動滑稽）？
廢話不多說，下面開始測試。

1.6 模型測試

在main.py中添加如下代碼：請忽略之前已經寫過的代碼show()和train()

import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
from dataset import Dataset
from models.MyCNN import MyCNN
from models.LeNet5 import LeNet5
import cv2

dst = Dataset()

def test():
  # 加載訓練好的模型
  model = torch.load('models/MyModels.pth')
  testing_correct = 0

  for batch, (imgs, labels) in enumerate(dst.loader_test, 1):
      if torch.cuda.is_available():
          X, y = imgs.cuda(), labels.cuda()
      else:
          X, y = imgs, labels
      outputs = model(X)
      _, pred = torch.max(outputs.detach(), dim=1)
      testing_correct += torch.sum(pred == y)
      if batch == len(dst.loader_test):
          print('Batch {}/{}, Test Acc:{}/{}={:.2f}%'.format(
              batch, len(dst.loader_test), testing_correct.item(),
              batch * dst.batch_size, 100 * testing_correct.item() / (batch * dst.batch_size)
          ))
if __name__ == '__main__':
  # train()
  test()

查看測試結果：

至此，我們的MNIST手寫數字識別已經大功告成了。

項目完整代碼請移步至我的github查看，如果覺得還不錯的話，希望能得到你的Star喲~