• Ai项目十四：基于 LeNet5 的手写数字识别及训练

    若该文为原创文章，转载请注明原文出处。

    一、介绍

    pytorch复现lenet5模型，并检测自己手写的数字图片。

    利用torch框架搭建模型相对比较简单，但是也会遇到很多问题，网上资料很多，搭建模型的方法大同小异，在我尝试了自己搭建搭建出来模型，无论是训练还是检测都会遇到很多的问题，像这种自己遇到的问题，请教别人也没有用。原本使用的是github上的一份代码来复现，环境搭建完成后，才发现要有GPU，而我搭建是使用CPU，失败告终，为了复现，租用了AutoDL平台，在次搭建，这里记录GPU下的操作，CPU版本需要修改源码，自行修改，我的目的是在要训练自己的模型并在RK3568上部署，所以先训练并测试好。为后续部署作基础。

    二、环境

    三、搭建

    1、创建虚拟环境

     conda create -n LeNet5_env python==3.8

    2、安装pytorch

    Previous PyTorch Versions | PyTorch

    根据官方PyTorch，安装pytorch，使用的是CPU版本，其他版本自行安装，安装命令：

    1. ​​​​​​​
    2. pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
    3.  -i https://pypi.tuna.tsinghua.edu.cn/simple

    还需要安装一些其他的库

    1. pip install matplotlib -i https://pypi.tuna.tsinghua.edu.cn/simple
    2. pip install opencv-python -i https://pypi.tuna.tsinghua.edu.cn/simple

    3、数据集下载

    1. http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
    2. http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
    3. http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
    4. http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz

    直接把上面地址复制到网页上，就只可以下载

    下载后保存到data/MNIST/raw目录下

    四、训练代码

    训练模型有四个文件分别为：LeNet5.py；myDatast.py；readMnist.py；train.py

    文件LeNet5.py是网络层模型

    train.py

    1. import torch
    2. from torch.autograd import Variable
    3. import torch.nn as nn
    4. from torch.utils.data import DataLoader
    5. from readMnist import *
    6. from myDatast import Mnist
    7. from LeNet5 import LeNet5
    8.  
    9. train_images = load_train_images()
    10. train_labels = load_train_labels()
    11.  
    12. trainData = Mnist(train_images, train_labels)
    13. train_data = DataLoader(dataset=trainData, batch_size=1, shuffle=True)
    14.  
    15. lenet5 = LeNet5()
    16. lenet5.cuda()
    17.  
    18. lossFun = nn.CrossEntropyLoss()
    19.  
    20. optimizer = torch.optim.Adam(params=lenet5.parameters(), lr=1e-4)
    21.  
    22. Epochs = 100
    23. L = len(train_data)
    24.  
    25. for epoch in range(Epochs):
    26.     for i, (img, id) in enumerate(train_data):
    27.  
    28.         img = img.float()
    29.         id = id.float()
    30.  
    31.         img = img.cuda()
    32.         id = id.cuda()
    33.  
    34.         img = Variable(img, requires_grad=True)
    35.         id = Variable(id, requires_grad=True)
    36.  
    37.         Output = lenet5.forward(img)
    38.         loss = lossFun(Output, id.long())
    39.  
    40.         optimizer.zero_grad()
    41.         loss.backward()
    42.         optimizer.step()
    43.  
    44.         iter = epoch * L + i + 1
    45.         if iter % 100 == 0:
    46.             print('epoch：{}，iter:{},loss:{:.6f}'.format(epoch + 1, iter, loss))
    47.  
    48.     torch.save(lenet5.state_dict(), 'lenet5.pth')
    49.

     LeNet5.py

    1. import torch.nn as nn
    2.  
    3.  
    4. class LeNet5(nn.Module):
    5.     def __init__(self):
    6.         super(LeNet5, self).__init__()
    7.  
    8.         self.conv1 = nn.Sequential(
    9.             nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5),
    10.             nn.Sigmoid(),
    11.             nn.MaxPool2d(kernel_size=2, stride=2)
    12.         )
    13.  
    14.         self.conv2 = nn.Sequential(
    15.             nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5),
    16.             nn.Sigmoid(),
    17.             nn.MaxPool2d(kernel_size=2, stride=2)
    18.         )
    19.  
    20.         self.fc1 = nn.Sequential(
    21.             nn.Linear(in_features=16 * 4 * 4, out_features=120),
    22.             nn.Sigmoid()
    23.         )
    24.  
    25.         self.fc2 = nn.Sequential(
    26.             nn.Linear(in_features=120, out_features=84),
    27.             nn.Sigmoid()
    28.         )
    29.  
    30.         self.fc3 = nn.Linear(in_features=84, out_features=10)
    31.  
    32.     def forward(self, img):
    33.         img = self.conv1.forward(img)
    34.         img = self.conv2.forward(img)
    35.  
    36.         img = img.view(img.size()[0], -1)
    37.  
    38.         img = self.fc1.forward(img)
    39.         img = self.fc2.forward(img)
    40.         img = self.fc3.forward(img)
    41.  
    42.         return img

     readMnist.py

    1. from torch.utils.data import Dataset
    2. from torchvision import transforms
    3. import numpy as np
    4.  
    5.  
    6. class Mnist(Dataset):
    7.     def __init__(self, dataset, label):
    8.         self.dataset = dataset
    9.         self.label = label
    10.         self.len = len(self.label)
    11.         self.transforms = transforms.Compose([transforms.ToTensor() , transforms.Normalize(mean=[0.5], std=[0.5])])
    12.  
    13.     def __len__(self):
    14.         return self.len
    15.  
    16.     def __getitem__(self, item):
    17.         img = self.dataset[item]
    18.         img_id = self.label[item]
    19.  
    20.         img = np.transpose(img,(1,2,0))
    21.         img = self.transforms(img)
    22.  
    23.         return img, img_id

    readMnist.py

    1. import numpy as np
    2. import struct
    3. import matplotlib.pyplot as plt
    4. import cv2
    5.  
    6. fpath = 'G:/enpei_Project_Code/21_LeNet5/LeNet5-master/myLeNet5/data/MNIST/raw/'
    7.  
    8. # 训练集文件
    9. train_images_idx3_ubyte_file = fpath + 'train-images-idx3-ubyte'
    10. # 训练集标签文件
    11. train_labels_idx1_ubyte_file = fpath + 'train-labels-idx1-ubyte'
    12.  
    13. # 测试集文件
    14. test_images_idx3_ubyte_file = fpath + 't10k-images-idx3-ubyte'
    15. # 测试集标签文件
    16. test_labels_idx1_ubyte_file = fpath + 't10k-labels-idx1-ubyte'
    17.  
    18.  
    19. def decode_idx3_ubyte(idx3_ubyte_file):
    20.     """
    21.     解析idx3文件的通用函数
    22.     :param idx3_ubyte_file: idx3文件路径
    23.     :return: 数据集
    24.     """
    25.     # 读取二进制数据
    26.     bin_data = open(idx3_ubyte_file, 'rb').read()
    27.  
    28.     # 解析文件头信息，依次为魔数、图片数量、每张图片高、每张图片宽
    29.     offset = 0
    30.     fmt_header = '>iiii'  # 因为数据结构中前4行的数据类型都是32位整型，所以采用i格式，但我们需要读取前4行数据，所以需要4个i。我们后面会看到标签集中，只使用2个ii。
    31.     magic_number, num_images, num_rows, num_cols = struct.unpack_from(fmt_header, bin_data, offset)
    32.     print('魔数:%d, 图片数量: %d张, 图片大小: %d*%d' % (magic_number, num_images, num_rows, num_cols))
    33.  
    34.     # 解析数据集
    35.     image_size = num_rows * num_cols
    36.     offset += struct.calcsize(fmt_header)  # 获得数据在缓存中的指针位置，从前面介绍的数据结构可以看出，读取了前4行之后，指针位置（即偏移位置offset）指向0016。
    37.     print(offset)
    38.     fmt_image = '>' + str(
    39.         image_size) + 'B'  # 图像数据像素值的类型为unsigned char型，对应的format格式为B。这里还有加上图像大小784，是为了读取784个B格式数据，如果没有则只会读取一个值（即一副图像中的一个像素值）
    40.     print(fmt_image, offset, struct.calcsize(fmt_image))
    41.     images = np.empty((num_images, 1, num_rows, num_cols))
    42.     # plt.figure()
    43.     for i in range(num_images):
    44.         if (i + 1) % 10000 == 0:
    45.             print('已解析 %d' % (i + 1) + '张')
    46.             print(offset)
    47.         images[i] = np.array(struct.unpack_from(fmt_image, bin_data, offset)).reshape((1, num_rows, num_cols))
    48.         # print(images[i])
    49.         offset += struct.calcsize(fmt_image)
    50.     #        plt.imshow(images[i],'gray')
    51.     #        plt.pause(0.00001)
    52.     #        plt.show()
    53.     # plt.show()
    54.  
    55.     return images
    56.  
    57.  
    58. def decode_idx1_ubyte(idx1_ubyte_file):
    59.     """
    60.     解析idx1文件的通用函数
    61.     :param idx1_ubyte_file: idx1文件路径
    62.     :return: 数据集
    63.     """
    64.     # 读取二进制数据
    65.     bin_data = open(idx1_ubyte_file, 'rb').read()
    66.  
    67.     # 解析文件头信息，依次为魔数和标签数
    68.     offset = 0
    69.     fmt_header = '>ii'
    70.     magic_number, num_images = struct.unpack_from(fmt_header, bin_data, offset)
    71.     print('魔数:%d, 图片数量: %d张' % (magic_number, num_images))
    72.  
    73.     # 解析数据集
    74.     offset += struct.calcsize(fmt_header)
    75.     fmt_image = '>B'
    76.     labels = np.empty(num_images)
    77.     for i in range(num_images):
    78.         if (i + 1) % 10000 == 0:
    79.             print('已解析 %d' % (i + 1) + '张')
    80.         labels[i] = struct.unpack_from(fmt_image, bin_data, offset)[0]
    81.         offset += struct.calcsize(fmt_image)
    82.     return labels
    83.  
    84.  
    85. def load_train_images(idx_ubyte_file=train_images_idx3_ubyte_file):
    86.     """
    87.     TRAINING SET IMAGE FILE (train-images-idx3-ubyte):
    88.     [offset] [type]          [value]          [description]
    89.     0000     32 bit integer  0x00000803(2051) magic number
    90.     0004     32 bit integer  60000            number of images
    91.     0008     32 bit integer  28               number of rows
    92.     0012     32 bit integer  28               number of columns
    93.     0016     unsigned byte   ??               pixel
    94.     0017     unsigned byte   ??               pixel
    95.     ........
    96.     xxxx     unsigned byte   ??               pixel
    97.     Pixels are organized row-wise. Pixel values are 0 to 255. 0 means background (white), 255 means foreground (black).
    99.     :param idx_ubyte_file: idx文件路径
    100.     :return: n*row*col维np.array对象，n为图片数量
    101.     """
    102.     return decode_idx3_ubyte(idx_ubyte_file)
    103.  
    104.  
    105. def load_train_labels(idx_ubyte_file=train_labels_idx1_ubyte_file):
    106.     """
    107.     TRAINING SET LABEL FILE (train-labels-idx1-ubyte):
    108.     [offset] [type]          [value]          [description]
    109.     0000     32 bit integer  0x00000801(2049) magic number (MSB first)
    110.     0004     32 bit integer  60000            number of items
    111.     0008     unsigned byte   ??               label
    112.     0009     unsigned byte   ??               label
    113.     ........
    114.     xxxx     unsigned byte   ??               label
    115.     The labels values are 0 to 9.
    117.     :param idx_ubyte_file: idx文件路径
    118.     :return: n*1维np.array对象，n为图片数量
    119.     """
    120.     return decode_idx1_ubyte(idx_ubyte_file)
    121.  
    122.  
    123. def load_test_images(idx_ubyte_file=test_images_idx3_ubyte_file):
    124.     """
    125.     TEST SET IMAGE FILE (t10k-images-idx3-ubyte):
    126.     [offset] [type]          [value]          [description]
    127.     0000     32 bit integer  0x00000803(2051) magic number
    128.     0004     32 bit integer  10000            number of images
    129.     0008     32 bit integer  28               number of rows
    130.     0012     32 bit integer  28               number of columns
    131.     0016     unsigned byte   ??               pixel
    132.     0017     unsigned byte   ??               pixel
    133.     ........
    134.     xxxx     unsigned byte   ??               pixel
    135.     Pixels are organized row-wise. Pixel values are 0 to 255. 0 means background (white), 255 means foreground (black).
    137.     :param idx_ubyte_file: idx文件路径
    138.     :return: n*row*col维np.array对象，n为图片数量
    139.     """
    140.     return decode_idx3_ubyte(idx_ubyte_file)
    141.  
    142.  
    143. def load_test_labels(idx_ubyte_file=test_labels_idx1_ubyte_file):
    144.     """
    145.     TEST SET LABEL FILE (t10k-labels-idx1-ubyte):
    146.     [offset] [type]          [value]          [description]
    147.     0000     32 bit integer  0x00000801(2049) magic number (MSB first)
    148.     0004     32 bit integer  10000            number of items
    149.     0008     unsigned byte   ??               label
    150.     0009     unsigned byte   ??               label
    151.     ........
    152.     xxxx     unsigned byte   ??               label
    153.     The labels values are 0 to 9.
    155.     :param idx_ubyte_file: idx文件路径
    156.     :return: n*1维np.array对象，n为图片数量
    157.     """
    158.     return decode_idx1_ubyte(idx_ubyte_file)
    159.  
    160.  
    161. if __name__ == '__main__':
    162.  
    163.     train_images = load_train_images()
    164.     train_labels = load_train_labels()
    165.     test_images = load_test_images()
    166.     test_labels = load_test_labels()
    167.  
    168.     pass
    169.  
    170.     # 查看前十个数据及其标签以读取是否正确
    171.     for i in range(10):
    172.         print(train_labels[i])
    173.  
    174.         img = train_images[i]
    175.         img = np.transpose(img, (1, 2, 0))
    176.  
    177.         cv2.namedWindow('img')
    178.         cv2.imshow('img', img)
    179.         cv2.waitKey(100)
    180.  
    181.     print('done')
    182.  
    183.

    上面代码需要注意的是数据集的路径，需要修改成对应的路径。

    运行python train.py

    训练大概5小时

    五、测试

    1. from LeNet5 import LeNet5
    2. import torch
    3. from readMnist import *
    4. from myDatast import Mnist
    5. from torch.utils.data import DataLoader
    6. import numpy as np
    7. import cv2
    8.  
    9. test_images = load_test_images()
    10. test_labels = load_test_labels()
    11.  
    12. testData = Mnist(test_images, test_labels)
    13. test_data = DataLoader(dataset=testData, batch_size=1, shuffle=True)
    14.  
    15. lenet5 = LeNet5()
    16. lenet5.load_state_dict(torch.load('lenet5.pth'))
    17. lenet5.eval()
    18.  
    19. showimg = True
    20. js = 0
    21. for i, (img, id) in enumerate(test_data):
    22.  
    23.     img = img.float()
    24.     outid = lenet5(img)
    25.  
    26.     oid = torch.argmax(outid)
    27.     if oid == id:
    28.         js = js + 1
    29.  
    30.     if showimg == True:
    31.         img = img.numpy()
    32.         img = np.squeeze(img)
    33.  
    34.         id = id.numpy()
    35.         id = np.squeeze(id)
    36.         id = np.int32(id)
    37.  
    38.         oid = oid.numpy()
    39.         oid = np.squeeze(oid)
    40.  
    41.         maxv = np.max(img)
    42.         minv = np.min(img)
    43.  
    44.         img = (img - minv) / (maxv - minv)
    45.  
    46.         cv2.namedWindow("img", 0)
    47.         cv2.imshow("img", img)
    48.  
    49.         title = "img, predicted value:{},truth value:{}".format(oid, id)
    50.         cv2.setWindowTitle("img",title)
    51.  
    52.         cv2.waitKey(1)
    53.  
    54. print('准确率：{:.6f}'.format(js / (i + 1)))

    测试结果准确率达到0.986基本达到要求 

    如有侵权，或需要完整代码，请及时联系博主。

  • 相关阅读:
    如何获取第三方maven依赖信息？
    云计算基础知识
    [vue] 嵌套iframe，$router.go(-1)后退bug
    iNFTnews | 元宇宙浪潮下，企业正通过AR和VR技术改善客户体验
    The Log-Structured Merge-Tree (LSM-Tree) 论文阅读笔记
    flutter 初视回味
    在c++项目中你必须真正使用的15个c++特性
    IDEA运行main方法，为什么要编译整个工程？
    Spring Data Jpa 多数据源配置（mysql 和 orcle）过程
    Linux 手动卸载jdk
  • 原文地址：https://blog.csdn.net/weixin_38807927/article/details/133514669