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

迈不过友情╰ 2023-07-12 06:20 58阅读 0赞

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

### 实验环境 ###

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次迭代

![20200307163750964.jpg][]

跑完了200次

![2020030716380530.jpg][]

### 完整代码 ###

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()

[20200307163750964.jpg]: https://img-blog.csdnimg.cn/20200307163750964.jpg
[2020030716380530.jpg]: https://img-blog.csdnimg.cn/2020030716380530.jpg