pytorch实现手写数字识别（LeNet-5模型）

LeNET-5模型如下图：

LeNET-5.png

模型中输入图片的大小为32x32，由于自身数据集的大小为784，即一张图片大小为28x28。所以稍微有所不同。数据集外网下载可能比较慢，将在文章最下方给出百度云下载地址。

下列文件分别为模型建立、模型参数存储、模型测试、模型训练。

文件结构

第一步：建立模型，即LeNET_5.py

import torch
import torch.nn as nn

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1=nn.Sequential(
            nn.Conv2d(in_channels=1,out_channels=6,kernel_size=5,stride=1),
            nn.MaxPool2d(kernel_size=2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1),
            nn.MaxPool2d(kernel_size=2)
        )
        self.fc1=nn.Sequential(
            nn.Linear(in_features=4*4*16,out_features=120)
        )
        self.fc2 = nn.Sequential(
            nn.Linear(in_features=120, out_features=84)
        )
        self.fc3 = nn.Sequential(
            nn.Linear(in_features=84, out_features=10)
        )

    def forward(self, input):
        conv1_output=self.conv1(input)#[28,28,1]--->[24,24,6]--->[12,12,6]
        conv2_output = self.conv2(conv1_output)  # [12,12,6]--->[8,8,,16]--->[4,4,16]
        conv2_output=conv2_output.view(-1,4*4*16)#将[n,4,4,16]维度转化为[n,4*4*16]
        fc1_output=self.fc1(conv2_output)#[n,256]--->[n,120]
        fc2_output=self.fc2(fc1_output)#[n,120]-->[n,84]
        fc3_output = self.fc3(fc2_output)  # [n,84]-->[n,10]
        return fc3_output

第二步：模型建立好后，进行训练，并保存模型参数，即LeNet_train.py

import torch
import torch.nn as nn
import torch.optim as optim
from classical_model.LeNET_5 import LeNet  #导入刚才建立的模型
import pandas as pd
import pandas.util
import matplotlib.pyplot as plt

train_data=pd.DataFrame(pd.read_csv("D:\program\py_workspace\data\mnist_csv\mnist_train.csv"))
model=LeNet()
print(model)
loss_fc=nn.CrossEntropyLoss() #定义损失函数
optimizer=optim.SGD(params=model.parameters(),lr=0.001) #采用随机梯度下降SGD
loss_list=[]#记录每次的损失值
x=[]#记录训练次数
for i in range(1000):
    batch_data=train_data.sample(n=30,replace=False)#每次随机读取30条数据
    batch_y=torch.from_numpy(batch_data.iloc[:,0].values).long()#标签值
    batch_x=torch.from_numpy(batch_data.iloc[:,1::].values).float().view(-1,1,28,28)
    #图片信息，一条数据784维将其转化为通道数为1，大小28*28的图片。

    prediction=model.forward(batch_x)#前向传播
    loss=loss_fc(prediction,batch_y)#计算损失值

    optimizer.zero_grad()
    loss.backward()#反向传播
    optimizer.step()#更新权重
    print("第%d次训练，loss为%.3f"%(i,loss))
    loss_list.append(loss)
    x.append(i)
torch.save(model.state_dict(),"../model/LeNet.pkl")#保存模型参数
plt.figure()
plt.plot(x,loss_list,"r-")#可以将损失值进行绘制
plt.show()

第三步：将训练好的模型用测试数据进行测试，即LeNet_test.py

import torch
import torch.nn as nn
from classical_model.LeNET_5 import LeNet
import pandas as pd
import numpy as np

model=LeNet()
test_data=pd.DataFrame(pd.read_csv("D:\program\py_workspace\data\mnist_csv\mnist_test.csv"))
model.load_state_dict(torch.load("../model/LeNet.pkl"))#加载模型参数

with torch.no_grad():#测试不需要反向传播
    for i in range(100):
        batch_data=test_data.sample(n=50,replace=False)
        batch_x=torch.from_numpy(batch_data.iloc[:,1::].values).float().view(-1,1,28,28)
        batch_y=batch_data.iloc[:,0].values
        prediction=np.argmax(model(batch_x).numpy(),axis=1)
        acccurcy=np.mean(prediction==batch_y)
        print("第%d组测试集，准确率为%.3f"%(i,acccurcy))

测试结果：

测试结果
下面给出数据集的下载地址
链接：https://pan.baidu.com/s/15GALKfgK2LLXBNit2lHXcg
提取码：ooca