Pytorch用LeNet在MNIST数据集上实现手写数字识别

本文主要讲解该算法的实现过程

实验环境

python3.6+pytorch1.2+cuda10.1

数据集

MNIST数据集包含6万训练图像与1万验证图像

下边是代码实现过程及讲解

数据加载

#载入数据
def loadMNIST(batch_size):
    trans_img = transforms.Compose([
        transforms.ToTensor()
    ])

    trainset = MNIST('./data',train=True,transform=trans_img,download=True)
    testset = MNIST('./data',train=False,transform=trans_img,download=True)

    trainloader = DataLoader(trainset,batch_size=batch_size,shuffle=True,num_workers=4)
    testloader = DataLoader(testset,batch_size=batch_size,shuffle=False,num_workers=4)

    return trainset,testset,trainloader,testloader

下载MNIST数据集到项目中，并构建数据集

定义网络结构

class Lenet(nn.Module):
    def __init__(self):
        super(Lenet,self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1,6,3,stride=1,padding=1),
            nn.MaxPool2d(2,2),
            nn.Conv2d(6,16,5,stride=1,padding=1),
            nn.MaxPool2d(2,2)

        )

        self.fc = nn.Sequential(
            nn.Linear(576,120),
            nn.Linear(120,84),
            nn.Linear(84,10)
        )

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

两个卷集层+两个最大池化层+三个全连接层，这里要记得卷积后要将数据展开成一列再带入全连接层

参数设定

learning_rate = 1e-3
batch_size = 100
epoches = 200

lenet = Lenet()
lenet.cuda()

trainset,testset,trainloader,testloader = loadMNIST(batch_size)

criterian = nn.CrossEntropyLoss()
optimizer = optim.SGD(lenet.parameters(),lr=learning_rate)

网络训练

for i in range(epoches):
    running_loss = 0.
    running_acc = 0.

    for (img,label) in trainloader:
        img = img.cuda()
        label = label.cuda()

        optimizer.zero_grad()

        output = lenet(img)
        loss = criterian(output,label)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _,predict = torch.max(output,1)
        correct_num = (predict==label).sum()
        running_acc += correct_num.item()


    running_loss /= len(trainset)
    running_acc /= len(trainset)
    print("[%d/%d] Loss: %.5f, Acc: %.2f"%(i+1,epoches,running_loss,running_acc*100),"%")

结果检验

def test():
    running_loss = 0.
    running_acc = 0.
    for (img,label) in testloader:
        img = img.cuda()
        label = label.cuda()
        output = lenet(img)
        _,predict = torch.max(output,1)
        correct_num = (predict==label).sum()
        running_acc += correct_num.item()
    running_acc /= len(testset)
    print("after train, Acc: %.2f"%(running_acc*100),"%")

结果展示

先跑了20次迭代

跑完了200次

完整代码

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms
from torch import optim
import  torch.nn as nn



#载入数据
def loadMNIST(batch_size):
    trans_img = transforms.Compose([
        transforms.ToTensor()
    ])

    trainset = MNIST('./data',train=True,transform=trans_img,download=True)
    testset = MNIST('./data',train=False,transform=trans_img,download=True)

    trainloader = DataLoader(trainset,batch_size=batch_size,shuffle=True,num_workers=4)
    testloader = DataLoader(testset,batch_size=batch_size,shuffle=False,num_workers=4)

    return trainset,testset,trainloader,testloader

class Lenet(nn.Module):
    def __init__(self):
        super(Lenet,self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1,6,3,stride=1,padding=1),
            nn.MaxPool2d(2,2),
            nn.Conv2d(6,16,5,stride=1,padding=1),
            nn.MaxPool2d(2,2)

        )

        self.fc = nn.Sequential(
            nn.Linear(576,120),
            nn.Linear(120,84),
            nn.Linear(84,10)
        )

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

learning_rate = 1e-3
batch_size = 100
epoches = 200

lenet = Lenet()
lenet.cuda()

trainset,testset,trainloader,testloader = loadMNIST(batch_size)

criterian = nn.CrossEntropyLoss()
optimizer = optim.SGD(lenet.parameters(),lr=learning_rate)



def test():
    running_loss = 0.
    running_acc = 0.
    for (img,label) in testloader:
        img = img.cuda()
        label = label.cuda()
        output = lenet(img)
        _,predict = torch.max(output,1)
        correct_num = (predict==label).sum()
        running_acc += correct_num.item()
    running_acc /= len(testset)
    print("after train, Acc: %.2f"%(running_acc*100),"%")

for i in range(epoches):
    running_loss = 0.
    running_acc = 0.

    for (img,label) in trainloader:
        img = img.cuda()
        label = label.cuda()

        optimizer.zero_grad()

        output = lenet(img)
        loss = criterian(output,label)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _,predict = torch.max(output,1)
        correct_num = (predict==label).sum()
        running_acc += correct_num.item()


    running_loss /= len(trainset)
    running_acc /= len(trainset)
    print("[%d/%d] Loss: %.5f, Acc: %.2f"%(i+1,epoches,running_loss,running_acc*100),"%")


    if (i+1)%20==0:
        test()