pytorch复现lenet5模型，并检测自己手写的数字图片_pytorch读取自己的手写数字数据集

写在前面：
利用框架搭建模型相对比较简单，但是也会遇到很多问题，网上资料很多，搭建模型的方法大同小异，但是在我尝试了用两种方法搭建出来模型，无论是训练还是检测都会遇到很多的问题，像这种自己遇到的问题，请教别人也没有用。大多数资料上面都是用torch自带的mnist下载数据集，这样自己就跳过了数据整理这个阶段，虽然能训练，并且出来很好的结果，但总归不是自己整理的数据，对于新手朋友来说，就很难图形化的查看自己手写的图片的预测结果。

这个是我的项目文件夹结构

第一步：用pytorch框架搭建lenet5模型

首先需要熟悉lenet5的结构，它是检测手写数字，手写数字图像是28*28大小的一个单通道图片，

第一层为卷积层+池化层：conv2d+maxpool

利用卷积，卷积核大小为5*5，padding为0，stride为1，卷积核数量为20
in_channel=1
out_channel=20
kernel_size=5
stride=1

self.conv1=nn.Conv2d(in_channels=1,out_channels=20,kernel_size=5,stride=1)
1

池化为最大池化

self.pool1=nn.MaxPool2d(2)
1

第二层为卷积层+池化层：conv2d+maxpool

和第一层一样，第二层的输入就是第一层的输出
in_channel=20
out_channel=50
kerne_size=5
stride=1

self.conv2=nn.Conv2d(in_channels=20,out_channels=50,kernel_size=5,stride=1)
1

第二层池化

self.pool2=nn.MaxPool2d(2)
1

第三层和第四层为全连接层Linear

in_features=800
为什么是800
第一层卷积后得到特征图的大小为（W-F+2P)/S+1=2424，数量20个
第一层池化后得到特征图大小为1212，数量20
第二层卷积后得到特征图的大小为（W-F+2P)/S+1=88，数量50
第二层池化后得到的特征图大小为44，数量50
把50个44的特征图reshape成一个1维数组，就是44*50=800

第三层全连接层输出维500，
最后一层输出为10

        self.fc1=nn.Linear(in_features=800,out_features=500,bias=True)
        self.relu1=nn.ReLU()

        self.fc2=nn.Linear(in_features=500,out_features=10,bias=True)
        self.relu2=nn.ReLU()
1
2
3
4
5

下面我贴出用pytorch搭建的lenet5网络模型

'''
首先介绍一下这个py文件，文件名为lenet5_v03,版本为v03，因为之前我已经用被的方法复现过两次lenet5算法，这次是第三次，就给这个文件命名为v03。
这个文件是定义一个类class，这个类定义一下LeNet5网络模型，这个网络模型是基于pytorch框架的。
这里面定义了两个模型，我认为是同样的模型的两种书写方式，
'''
import torch
from torch import nn
from torch.nn import functional as F

class LeNet5(nn.Module):
    def __init__(self):
        super().__init__()
        self.cnn_layers=nn.Sequential(
             #定义卷积层，1个输入通道，6个输出通道，5*5的卷积filter
            nn.Conv2d(in_channels=1,out_channels=20,kernel_size=5,stride=1),
            nn.MaxPool2d(2),
            #the second cnn_layer,input 20 feature map,output 50 feature map,kernel_size=5,stride=1
            nn.Conv2d(in_channels=20,out_channels=50,kernel_size=5,stride=1),
            nn.MaxPool2d(2)
         )
        self.fc_layers=nn.Sequential(
            # 3 full connect layers
            nn.Linear(800,500),
            nn.ReLU(),
            nn.Linear(500,10),
            nn.LogSoftmax(dim=1)
        )
    def forward(self,x):
        #the forward function
        out=self.cnn_layers(x)
        out=out.view(-1,800)
        out=self.fc_layers(out)
        return out

class LeNet5_01(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1=nn.Conv2d(in_channels=1,out_channels=20,kernel_size=5,stride=1)
        self.pool1=nn.MaxPool2d(2)
        #the second conv input 20,output 50 kernel_size=5,stride=1
        self.conv2=nn.Conv2d(in_channels=20,out_channels=50,kernel_size=5,stride=1)
        self.pool2=nn.MaxPool2d(2)
        # the full connect
        self.fc1=nn.Linear(in_features=800,out_features=500,bias=True)
        self.relu1=nn.ReLU()

        self.fc2=nn.Linear(in_features=500,out_features=10,bias=True)
        self.relu2=nn.ReLU()

    def forward(self,x):
        x=self.conv1(x)
        x=self.pool1(x)
        x=self.conv2(x)
        x=self.pool2(x)
        x=x.view(-1,800)
        x=self.fc1(x)
        x=self.relu1(x)
        x=self.fc2(x)
        x=self.relu2(x)
        x=F.log_softmax(input=x,dim=1)
        return x



1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

上面的代码用了两种方式搭建了模型，分别是LeNet5和LeNet5_01

第二步：数据集整理

这里我不用pytorch自己下载数据集，这里我从网上下载mnist数据集
首先介绍一下mnist数据集

这里是我下载的mnist数据集的二进制文件，总共是4个文件，分别是训练集，训练集标签，测试集，测试集标签，其中训练集为60000张图片，测试集为10000张图片，同时对应的标签为60000个和10000个
图片信息是以二进制的格式保存在这四个文件中的，下面我们需要创建方法读取二进制文件中的内容

'''
解码二进制文件的函数
用来解码MNIST数据集里面的二进制文件

'''
import struct#struct模块
import numpy as np #numpy包

def decode_idx3_ubyte(idx3_ubyte_file):
    """
    解析idx3文件的通用函数
    :param idx3_ubyte_file: idx3文件路径
    :return: 数据集
    """
    # 读取二进制数据
    with open(idx3_ubyte_file, 'rb') as bin_data1:
        bin_data=bin_data1.read()
        # 解析文件头信息，依次为魔数、图片数量、每张图片高、每张图片宽
        offset = 0
        fmt_header = '>4i'
        '''使用大端法'''
        #mnist使用的大端方法存储的数据
        # 因为数据结构中前4行的数据类型都是32位整型，所以采用i格式，但我们需要读取前4行数据，所以需要4个i。我们后面会看到标签集中，只使用2个ii。
        magic_number, num_images, num_rows, num_cols = struct.unpack_from(fmt_header, bin_data, offset)
        print('魔数:%d, 图片数量: %d张, 图片大小: %d*%d' % (magic_number, num_images, num_rows, num_cols))

        # 解析数据集
        image_size = num_rows * num_cols
        # 获得数据在缓存中的指针位置，从前面介绍的数据结构可以看出，读取了前4行之后，指针位置（即偏移位置offset）指向0016。
        print(struct.calcsize(">4i"))
        offset =offset+ struct.calcsize(fmt_header)
        print(offset)
        # 图像数据像素值的类型为unsigned char型，对应的format格式为B。这里还有加上图像大小784，是为了读取784个B格式数据，如果没有则只会读取一个值（即一副图像中的一个像素值）
        #B是一个字节8为，I是4个字节32位
        fmt_image = '>' + str(image_size) + 'B'
        print(fmt_image,offset,struct.calcsize(fmt_image))
        #1万张图片
        images = np.empty((num_images, num_rows, num_cols))
        #plt.figure()
        #j=0
        for i in range(num_images):
            if (i + 1) % 10000 == 0:
                print('已解析 %d' % (i + 1) + '张')
                print(offset)
            #读取数据放入第i行，并reshape（28，28）
            images[i] = np.array(struct.unpack_from(fmt_image, bin_data, offset)).reshape((num_rows, num_cols))
            #print("输出",images[i])
            offset += struct.calcsize(fmt_image)
            #plt.imshow(images[i],'gray')
            #不明白是什么意思

            #plt.pause(0.001)
            #plt.show()
        #plt.show()

        return images

def decode_idx1_ubyte(idx1_ubyte_file):
    """
    解析idx1文件的通用函数
    :param idx1_ubyte_file: idx1文件路径
    :return: 数据集
    """
    # 读取二进制数据
    with open(idx1_ubyte_file, 'rb') as bin_data1:
        bin_data=bin_data1.read()
        # 解析文件头信息，依次为魔数、图片数量、每张图片高、每张图片宽
        offset = 0
        fmt_header = '>2i'
        '''使用大端法'''
        #mnist使用的大端方法存储的数据
        # 因为数据结构中前4行的数据类型都是32位整型，所以采用i格式，但我们需要读取前4行数据，所以需要4个i。我们后面会看到标签集中，只使用2个ii。
        magic_number, num_labels = struct.unpack_from(fmt_header, bin_data, offset)
        print('魔数:%d, 图片标签数量: %d个' % (magic_number, num_labels))

        # 解析数据集
        label_size = 1
        # 获得数据在缓存中的指针位置，从前面介绍的数据结构可以看出，读取了前4行之后，指针位置（即偏移位置offset）指向0016。
        print(struct.calcsize(">2i"))
        offset =offset+ struct.calcsize(fmt_header)
        print(offset)
        # 图像数据像素值的类型为unsigned char型，对应的format格式为B。这里还有加上图像大小784，是为了读取784个B格式数据，如果没有则只会读取一个值（即一副图像中的一个像素值）
        #B是一个字节8为，I是4个字节32位
        fmt_label = '>' + str(label_size) + 'B'
        print(fmt_label,offset,struct.calcsize(fmt_label))
        #1万张图片
        labels = np.empty((num_labels, 1))
        #plt.figure()
        j=0
        for i in range(num_labels):
            labels[i] = np.array(struct.unpack_from(fmt_label, bin_data, offset)).reshape(1)
            #print("输出",images[i])
            offset += struct.calcsize(fmt_label)
            #plt.imshow(images[i],'gray')
            #不明白是什么意思

            #plt.pause(0.001)
            #plt.show()
        #plt.show()

        return labels
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

上面的代码就是两个方法，一个是读取图片，并转换成ndarray数组，一个是读取labels，转换成ndarray数组，并返回数据

第三步：训练方法train(),训练数据

下面是文件train.py的代码

#author:chenchen
import torch as t
import numpy as np
from lenet5_v03 import LeNet5_01,LeNet5
from torch.utils.data import DataLoader,TensorDataset
from decode_binary_function import decode_idx3_ubyte,decode_idx1_ubyte
#定义一个train方法，训练模型
def train(EPOCH,model,train_dl):
    model.train()
    print('_'*10,"训练开始",'_'*10)
    print("model's state_dict:")
    for param_tensor in model.state_dict():
        print(param_tensor,"\t",model.state_dict()[param_tensor].size())
    loss=t.nn.CrossEntropyLoss()
    opt=t.optim.Adam(model.parameters(),lr=1e-3)
    for e in range(EPOCH):
        print("run in EPOCH:%d"%e)
        for i,(x_train,y_train) in enumerate(train_dl):
            x_train=x_train.cuda()
            y_train=y_train.cuda()
            y_pred=model.forward(x_train)
            train_loss=loss(y_pred,y_train)
            if (i+1)%100==0:
                print('batch:',i+1,train_loss.item())
                opt.zero_grad()
                train_loss.backward()
                opt.step()
    t.save(model.state_dict(),'wb.pt')
    print('*'*10,'训练完毕','*'*10)

#主程序
if __name__=="__main__":

    print('*' * 10,'程序开始执行......','*'*10)
    EPOCH = 50
    batch_size=32
    train_images_path=r"data/train-images-idx3-ubyte"
    train_labels_path=r"data/train-labels-idx1-ubyte"
    train_images=decode_idx3_ubyte(train_images_path)
    train_labels=decode_idx1_ubyte(train_labels_path)
    train_images=train_images.reshape(60000,1,28,28).astype(np.float32)/255-0.5/0.5
    train_labels=train_labels.reshape(60000).astype(np.long)
    train_images=t.from_numpy(train_images)
    train_labels=t.from_numpy(train_labels).type(t.long)
    #print(train_images[0])

    train_ds=TensorDataset(train_images,train_labels)
    model = LeNet5().cuda()
    train_dl=DataLoader(dataset=train_ds,batch_size=batch_size,shuffle=True,drop_last=False)
    train(EPOCH, model,train_dl)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

这里的代码主要定义了一个train方法，主程序就是加载数据，把数据转换成tensor，合并train和label为dataset，然后加载数据，
整个过程中用到的torch自带的数据集处理工具
重要应该理解的就是TensorDataset,DataLoader,这两个方法，代码我都是一步一步写的，可以很清晰的了解这两个方法的使用。
最后保存训练好的权重文件为wb.pt

第四步 test（），计算准确率

下面的就是test.py 这个文件的内容

#author=chenchen
import numpy as np
import torch as t
from torch.utils.data import TensorDataset,DataLoader
from lenet5_v03 import LeNet5,LeNet5_01
from decode_binary_function import decode_idx1_ubyte,decode_idx3_ubyte
#定义一个测试方法
def test(model,test_dl,wt):
    print("测试开始：")
    total=0
    correct_count=0
    model.eval()
    model.load_state_dict(t.load(wt))
    for i,(x_test,y_test) in enumerate(test_dl):
        pred_labels=model(x_test.cuda())
        predicted=t.max(pred_labels,1)[1]
        correct_count=correct_count+(predicted==y_test.cuda()).sum()
        total=total+len(y_test)
    print('total acc:%.2f\n'%(correct_count/total))

if __name__=="__main__":
    model = LeNet5().cuda()
    test_images_path = r"data/t10k-images-idx3-ubyte"
    test_labels_path = r"data/t10k-labels-idx1-ubyte"
    test_images = decode_idx3_ubyte(test_images_path)
    test_labels = decode_idx1_ubyte(test_labels_path)
    test_images = test_images.reshape(10000, 1, 28, 28).astype(np.float32) / 255 - 0.5 / 0.5
    test_labels = test_labels.reshape(10000).astype(np.long)
    test_images = t.from_numpy(test_images)
    test_labels = t.from_numpy(test_labels).type(t.long)
    test_ds = TensorDataset(test_images, test_labels)
    test_dl = DataLoader(dataset=test_ds, batch_size=62, shuffle=True)

    wt = "wb.pt"

    test(model=model, test_dl=test_dl, wt=wt)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

这个和训练类似，也是需要加载数据，构造加载器，加载模型，加载好训练好的权重文件，然后输入test的数据集，计算预测值
这里用的一个方法就是加载训练好的权重文件

第五步：detect（），可视化测试模型

这个文件名为detect.py

#author=chenchen
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
from lenet5_v03 import LeNet5
import torch as t
import cv2
def detect(model,image):
    print("预测开始：")
    model.eval()
    wt='wb.pt'
    model.load_state_dict(t.load(wt))
    image=t.from_numpy(image)
    pred_labels=model(image.cuda())
    predicted=t.max(pred_labels,1)[1].cpu()
    print(type(predicted))
    print(predicted.shape)
    num=predicted.numpy()
    print("num:",num[0])
def load_image(image_path):
    image=Image.open(image_path)
    plt.imshow(image)
    plt.show()
    image = np.array(image)
    image=image[:,:,0]
    a=image[0][0]-22
    print(a)
    print(image)
    image=Image.fromarray(image)
    #image=image.convert('L')
    plt.imshow(image)
    plt.show()
    #image.show()
    threshold=a
    table=[]
    for i in range(256):
        if i<threshold:
            table.append(1)
        else:
            table.append(0)
    image=image.point(table,"1")
    plt.imshow(image)
    plt.show()
    image=image.convert('L')
    image = image.resize((28, 28), Image.ANTIALIAS)
    plt.imshow(image)
    plt.show()
    image=np.array(image).reshape(1,1,28,28).astype('float32')
    image=image/255-0.5/0.5
    print(image)
    return image
def load_image1(file):
    img=cv2.imread(file)
    cv2.imshow("加载完成",img)
    cv2.waitKey(0)
    b,g,r=cv2.split(img)
    cv2.imshow("r",r)
    cv2.waitKey(0)
    threshold =100
    table = []
    for i in range(256):
        if i < threshold:
            table.append(1)
        else:
            table.append(0)

    # 图片二值化
    img=Image.fromarray(r)
    img = img.point(table, '1')
    plt.imshow(img)
    plt.show()
    print(type(img))
    img = img.convert('L')

    # 预处理
    # 调整图像大小
    plt.imshow(img)
    plt.show()

    img = img.resize((28,28),Image.ANTIALIAS)


    plt.imshow(img)
    plt.show()
    img = np.array(img).reshape(1,1,28,28).astype('float32')
    # 归一化处理
    img = img / 255-0.5/0.5
    return img

if __name__=="__main__":
    model=LeNet5().cuda()
    image_path = r"./detect_images/555.png"
    image=load_image(image_path)
    detect(model=model,image=image)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

这一个也是和train和test类似，主要是读取图片，处理图片，把图片转为tensor，归一化处理

这是我手写的数字，用手机拍的照片，在detect_images文件夹下面，由于拍摄环境的影响，我在前处理这些图片的时候需要把背景全部转换成黑色，数字转成白色，因为拍摄图片光照的原因，在0-256之间这些图片对应的一个阈值不一样，有差别，我就在每次读取图片的时候，读取一个大概的背景值对阈值

这里我读取的图片的第一个数值减去22为背景阈值，
下面我贴出几张图片
原始图像

单通道图像

背景黑色，数字白色处理

resize并归一化处理

预测结果

还有需要注意的是下面这个截图,这里又执行了一遍image.convert(‘L’)具体没有搞明白，转为8位像素