Pytorch實現手寫數字識別

LeNet-5網路的簡單實現

# 匯入相應的包
import torch.nn as nn

# 搭建LeNet-5
class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5,self).__init__()
        
        # 搭建卷積池化卷積池化層
        self.conv1 = nn.Sequential(
            nn.Conv2d(  # (32*32*1)
                in_channels = 1,
                out_channels = 6,
                kernel_size = 5,
                stride = 1,
                padding = 0
            ),  # (28*28*6)
            # 卷積和池化層之間新增ReLU啟用函式（原文使用Sigmoid或者tanh函式）
            nn.ReLU(),  
            nn.AvgPool2d(kernel_size = 2)  # (14*14*6)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(6,16,5,1,0),  # (10*10*16)
            nn.ReLU(),
            nn.AvgPool2d(2)  # (5*5*16)
        )
        
        # 搭建全連線層
        self.out = nn.Sequential(
            nn.Linear(16*5*5,120),
            nn.ReLU(),  # 全連線層之間使用ReLU啟用函式來進行啟用
            nn.Linear(120,84),
            nn.ReLU(),
            nn.Linear(84,10)
        )
        
# 搭建前向傳播介面
def forward(self,x):
    x = self.conv1(x)
    x = self.conv2(x)
    # 將卷積層最後一個輸出化成一個一維向量，然後才能傳入全連線層
    # 可以使用numpy中的reshape函式來進行替代
    x = x.view(x.size(0),-1)
    output = self.out(x)
    return output
    
# look your network
myNet = LeNet5()
print(myNet)

LeNet-5應用於手寫數字識別

Step1.匯入所需的包

# 導包
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets,transforms
from torch.autograd import Variable

Step2.確定引數值

# 引數選取
lr = 0.01
momentum = 0.5
log_interval = 10  # 跑多少次batch進行一次日誌記錄
epochs = 10
batch_size = 64
test_batch_size = 1000

Step3.構建LeNet5網路模型

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet,self).__init__()
        self.conv1 = nn.Sequential(  # (28*28*1)
            nn.Conv2d(1,6,5,1,2),  # (28*28*6)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 2,stride = 2)  # (14*14*6)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(6,16,5),
            nn.ReLU(),  # (10*10*16)
            nn.MaxPool2d(2,2)  # (5*5*16)
        )
        self.fc1 = nn.Sequential(
            nn.Linear(16*5*5,120),
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(120,84),
            nn.ReLU()
        )
        self.fc3 = nn.Linear(84,10)
    def forward(self,x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size()[0],-1)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

Step4.定義每次訓練的細節

def train(epoch):  # 定義每個epoch的訓練細節
    model.train()  # 設定為trainning模式
    for batch_idx, (data, target) in enumerate(train_loader):
        data = data.to(device)
        target = target.to(device)
        data, target = Variable(data), Variable(target)  # 把資料轉換成Variable
        optimizer.zero_grad()  # 優化器梯度初始化為零
        output = model(data)  # 把資料輸入網路並得到輸出，即進行前向傳播
        loss = F.cross_entropy(output,target)  #交叉熵損失函式
        loss.backward()  # 反向傳播梯度
        optimizer.step()  # 結束一次前傳+反傳之後，更新引數
        if batch_idx % log_interval == 0:  # 準備列印相關資訊
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))

Step5.定義測試函式

# 定義測試函式
def test():
    model.eval()  # 設定為test模式
    test_loss = 0  # 初始化測試損失值為0
    correct = 0  # 初始化預測正確的資料個數為0
    for data, target in test_loader:
 
        data = data.to(device)
        target = target.to(device)
        data, target = Variable(data), Variable(target)  #計算前要把變數變成Variable形式，因為這樣子才有梯度
 
        output = model(data)
        test_loss += F.cross_entropy(output, target, size_average=False).item()  # sum up batch loss 把所有loss值進行累加
        pred = output.data.max(1, keepdim=True)[1]  # get the index of the max log-probability
        correct += pred.eq(target.data.view_as(pred)).cpu().sum()  # 對預測正確的資料個數進行累加
 
    test_loss /= len(test_loader.dataset)  # 因為把所有loss值進行過累加，所以最後要除以總得資料長度才得平均loss
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

Step6.訓練預測模型

# 主函式啟動訓練
if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #啟用GPU
 

    train_loader = torch.utils.data.DataLoader(  # 載入訓練資料
    datasets.MNIST('../data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))  #資料集給出的均值和標準差係數，每個資料集都不同的，都資料集提供方給出的
                   ])),
    batch_size=batch_size, shuffle=True)

    test_loader = torch.utils.data.DataLoader(  # 載入訓練資料，詳細用法參考我的Pytorch打怪路（一）系列-（1）
    datasets.MNIST('../data', train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)) #資料集給出的均值和標準差係數，每個資料集都不同的，都資料集提供方給出的
    ])),
    batch_size=test_batch_size, shuffle=True)
    
 
    model = LeNet()  # 例項化一個網路物件
    model = model.to(device)
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)  # 初始化優化器
 
    for epoch in range(1, epochs + 1):  # 以epoch為單位進行迴圈
        train(epoch)
        test()
 
    torch.save(model, 'model.pth') #儲存模型

Step7.手寫數字識別測試
【注】使用畫圖板確保圖片畫素是28*28大小的，否則會報錯

# 傳入一張自己手寫的數字進行識別
# 可將此主函式放入另一個python檔案中
import torch
import cv2
import torch.nn.functional as F
# from modela import LeNet  ##重要，雖然顯示灰色(即在次程式碼中沒用到)，但若沒有引入這個模型程式碼，載入模型時會找不到模型
from torch.autograd import Variable
from torchvision import datasets, transforms
import numpy as np
 
if __name__ =='__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = torch.load('model.pth') #載入模型
    model = model.to(device)
    model.eval()    #把模型轉為test模式
 
    img = cv2.imread("E:/6.jpg")  #讀取要預測的圖片
    trans = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
 
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)#圖片轉為灰度圖，因為mnist資料集都是灰度圖
    img = trans(img)
    img = img.to(device)
    img = img.unsqueeze(0)  #圖片擴充套件多一維,因為輸入到儲存的模型中是4維的[batch_size,通道,長，寬]，而普通圖片只有三維，[通道,長，寬]
    #擴充套件後，為[1，1，28，28]
    output = model(img)
    prob = F.softmax(output, dim=1)
    prob = Variable(prob)
    prob = prob.cpu().numpy()  #用GPU的資料訓練的模型儲存的引數都是gpu形式的，要顯示則先要轉回cpu，再轉回numpy模式
    print(prob)  #prob是10個分類的概率
    pred = np.argmax(prob) #選出概率最大的一個
    print(pred.item())

輸入
預測輸出

【注】該模型的預測準確度還是有待提升，待後續引入更好的訓練模型進行訓練