用PyTorch搭建LeNet-5手写识别

首先申明，这篇博客用于记录本人看完LeNet-5论文，并对其中的算法进行复现的记录，可以看成是学习笔记

这里只介绍复现的工作，如果想了解更多有关网络的细节，请去看论文《Gradient-Based Learning Applied to Document Recognition》

在此推荐一个b站up的视频从0开始撸代码--手把手教你搭建LeNet-5网络模型_哔哩哔哩_bilibili，博主也是根据此视频进行复现的，博主其实是个小菜鸟

博主觉得up讲的还不错的，视频不涉及原理，只是手把手教你如何搭建

要想细追原理，最好直接看YannLeCun论文《Gradient BasedLearning Applied to Document Recognition》，在此奉上。

链接：https://pan.baidu.com/s/1cB1pheefesy2Q6aR2WscXg?pwd=iq43 提取码：iq43

一、必要的环境

如果你什么都不会，可以先去这篇博客把所需的驱动，软件都下好，里面paddlepaddle环境不用安装

这里博主也是重新创建了一个叫pytorch的环境，python版本是3.8，

然后在cmd输入nvidia-smi命令来查看自己电脑最高支持的cuda版本

我的最高支持是11.7，我下载的是cuda11.3版本的

在之前创建的pytoch输入代码，应该就能安装成功

conda install pytorch torchvision torchaudio cudatoolkit=11.3

但博主输入这行代码就会报错，好像是找不到库还是什么原因，如果你们也会报错试试下面的代码

conda install pytorch==1.11.0 torchvision=0.12.0 cudatoolkit=11.3 -c pytorch

后面的-C不能去掉，这样下载能成功，但速度有点慢

注意：里面一个pytorch包1.2g太大了，如果因为网速慢没下载成功，可以试试这串代码

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/

然后再次输入这行代码conda install pytorch==1.11.0 torchvision=0.12.0 cudatoolkit=11.3 -c pytorch ，把剩下的包下载好

下载好后，用炮哥博客的代码进行验证

import torch

print(torch.cuda.is_available())

print(torch.backends.cudnn.is_available())

print(torch.cuda_version)

print(torch.backends.cudnn.version())

结果显示，就表示成功了，cuda版本11.3，cudnn的版本为8.20版本

到此为止，手写识别所需的环境就安装好了

二、搭建模型、训练

1.整体框图

我们就要利用整体框图来搭建模型，卷积层都采用5*5的卷积核，步长为1，池化层(下采样层)采用都2*2的卷积，步长为2

2.net.py

搭建模型基本结构、手写识别的代码还是比较好看懂的，可以自己去理解下

 1 import torch

 2 from torch import nn

 3

 4 #定义一个网络模型类

 5 class MyLeNet5(nn.Module):

 6     #初始化网络

 7     def __init__(self):

 8         super(MyLeNet5,self).__init__()

 9         #输入大小为32*32，输出大小为28*28,输入通道为1，输出为6，卷积核大小为5,步长为1

10         self.c1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, padding=2)

11         #sigmoid激活函数

12         self.Sigmoid= nn.Sigmoid()

13         #平均池化

14         self.s2 = nn.AvgPool2d(kernel_size=2, stride=2)

15         self.c3 = nn.Conv2d(in_channels=6,out_channels=16,kernel_size=5)

16         self.s4 = nn.AvgPool2d(kernel_size=2,stride=2)

17         self.c5 = nn.Conv2d(in_channels=16,out_channels=120,kernel_size=5)

18         #展开

19         self.flatten = nn.Flatten()

20         self.f6 = nn.Linear(120,84)

21         self.output = nn.Linear(84,10)

22

23     def forward(self,x):

24         #输入x为32*32*1,输出为28*28*6

25         x = self.Sigmoid(self.c1(x))

26         #输入为28*28*6,输出为14*14*6

27         x = self.s2(x)

28         # 输入为14*14*6,输出为10*10*16

29         x = self.Sigmoid(self.c3(x))

30         # 输入为10*10*16,输出为5*5*16

31         x = self.s4(x)

32         # 输入为5*5*16,输出为1*1*120

33         x = self.c5(x)

34         x = self.flatten(x)

35         # 输入为120，输出为84

36         x = self.f6(x)

37         # 输入为84，输出为10

38         x = self.output(x)

39         return x

40

41 if __name__=="__main__":

42     x = torch.rand([1,1,28,28])#任意产生一个张量,批次1，通道为1，大小为28*28

43     model = MyLeNet5()#网络实例化

44     y = model(x) #输出结果

写完后保存，可以运行下看是否报错

3.train.py

这是用于训练模型的代码

  1 import torch

  2 from torch import nn

  3 from net import MyLeNet5

  4 from torch.optim import  lr_scheduler

  5 from torchvision import datasets,transforms

  6 import os

  7

  8

  9 #将数据转化为tensor格式

 10 data_transform = transforms.Compose([

 11     transforms.ToTensor()

 12 ])

 13

 14 # 加载训练数据集

 15 train_dataset = datasets.MNIST(root='./data', train=True, transform=data_transform, download=True)

 16 # 给训练集创建一个数据加载器, shuffle=True用于打乱数据集，每次都会以不同的顺序返回。

 17 train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)

 18 # 加载训练数据集

 19 test_dataset = datasets.MNIST(root='./data', train=False, transform=data_transform, download=True)

 20 # 给训练集创建一个数据加载器, shuffle=True用于打乱数据集，每次都会以不同的顺序返回。

 21 test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=16, shuffle=True)

 22

 23

 24 # 如果显卡可用，则用显卡进行训练

 25 device = "cuda" if torch.cuda.is_available() else 'cpu'

 26

 27 #调用net文件的模型，果GPU可用则将模型转到GPU

 28 model = MyLeNet5().to(device)

 29

 30 #定义损失函数，交叉熵损失

 31 loss_fn = nn.CrossEntropyLoss()

 32

 33 #定义优化器SGD,随机梯度下降

 34 optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)

 35

 36 #学习率每10个epoch变为原来的0.1

 37 lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

 38

 39 #定义训练函数

 40 def train(dataloader, model, loss_fn, optimizer):

 41     loss, current, n = 0.0, 0.0, 0

 42     # enumerate返回为数据和标签还有批次

 43     for batch, (X, y) in enumerate(dataloader):

 44         # 前向传播

 45         X, y = X.to(device), y.to(device)

 46         output = model(X)

 47         cur_loss = loss_fn(output, y)

 48         # torch.max返回每行最大的概率和最大概率的索引,由于批次是16，所以返回16个概率和索引

 49         _, pred = torch.max(output, axis=1)

 50

 51         # 计算每批次的准确率， output.shape[0]为该批次的多少

 52         cur_acc = torch.sum(y == pred) / output.shape[0]

 53         # print(cur_acc)

 54         # 反向传播

 55         optimizer.zero_grad()

 56         cur_loss.backward()

 57         optimizer.step()

 58         # 取出loss值和精度值

 59         loss += cur_loss.item()

 60         current += cur_acc.item()

 61         n = n + 1

 62

 63     print('train_loss:' + str(loss / n))

 64     print('train_acc:' + str(current / n))

 65

 66

 67 #定义验证函数

 68 def val(dataloader,model,loss_fn):

 69     # 将模型转为验证模式

 70     model.eval()

 71     loss, acc, n = 0.0, 0.0, 0

 72     # enumerate返回为数据和标签还有批次

 73     with torch.no_grad():

 74         for batch, (x, y) in enumerate(dataloader):

 75             # 前向传播

 76             x, y = x.to(device), y.to(device)

 77             output = model(x)

 78             cur_loss = loss_fn(output, y)

 79             # torch.max返回每行最大的概率和最大概率的索引,由于批次是16，所以返回16个概率和索引

 80             _, pred = torch.max(output, axis=1)

 81

 82             # 计算每批次的准确率， output.shape[0]为该批次的多少

 83             cur_acc = torch.sum(y == pred) / output.shape[0]

 84             loss += cur_loss.item()

 85             acc += cur_acc.item()#取出单元素张量的元素值并返回该值

 86             n += 1  # 记录有多少批次

 87         print('test_loss:' + str(loss / n))

 88         print('test_acc:' + str(acc / n))

 89

 90         return acc/n

 91

 92 #开始训练

 93 epoch = 30#训练轮次

 94 max_acc = 0

 95 for t in range(epoch):

 96     lr_scheduler.step()#学习率调整

 97     print(f"epoch{t+1}\n-------------------")#加f表示格式化字符串，加f后可以在字符串里面使用用花括号括起来的变量和表达式

 98     train(train_dataloader, model, loss_fn, optimizer)#调用train函数

 99     a = val(test_dataloader,model,loss_fn)

100     #保存最后的模型权重文件

101     if a > max_acc:

102         folder = 'save_model'

103         if not os.path.exists(folder):

104             os.mkdir('save_model')

105         max_acc = a

106         print('save best model')

107         torch.save(model.state_dict(),"save_model/best_model.pth")

108     #保存最后的文件

109     if t == epoch - 1:

110         torch.save(model.state_dict(),"save_model/last_model.pth")

111 print('Done')

写完后运行train.py，大概需要一会时间，代码运行完成后，会生成最好和最后的权重

博主训练了30轮，训练集和测试集的准确就达到了96

三、模型测试

1.test.py

训练完成后，将最好的权重路径放到test.py文件里，运行代码,在此博客选择前10张图片作为验证，你们可以根据需求自己改

 1 import torch

 2 from net import MyLeNet5

 3 from torch.autograd import Variable

 4 from torchvision import datasets,transforms

 5 from torchvision.transforms import ToPILImage

 6

 7 # 将数据转化为tensor格式

 8 data_transform = transforms.Compose([

 9     transforms.ToTensor()

10 ])

11

12 # 加载训练数据集

13 train_dataset = datasets.MNIST(root='./data', train=True, transform=data_transform, download=True)

14 # 给训练集创建一个数据加载器, shuffle=True用于打乱数据集，每次都会以不同的顺序返回。

15 #train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)

16 # 加载训练数据集

17 test_dataset = datasets.MNIST(root='./data', train=False, transform=data_transform, download=True)

18 # 给训练集创建一个数据加载器, shuffle=True用于打乱数据集，每次都会以不同的顺序返回。

19 #test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=16, shuffle=True)

20

21 #  如果显卡可用，则用显卡进行训练

22 device = "cuda" if torch.cuda.is_available() else 'cpu'

23

24 # 调用net里面定义的模型，如果GPU可用则将模型转到GPU

25 model = MyLeNet5().to(device)

26

27 #加载train.py里训练好的模型

28 model.load_state_dict(torch.load(("D:/python/LeNet-5/save_model/best_model.pth")))#填写权重路径

29

30 #获取预测结果

31

32 classes = [

33     "0",

34     "1",

35     "2",

36     "3",

37     "4",

38     "5",

39     "6",

40     "7",

41     "8",

42     "9",

43 ]

44

45 # 把tensor转换成Image，方便可视化

46 show = ToPILImage()

47

48 #进入验证阶段

49 model.eval()

50 # 对test_dataset手写数字图片进行推理

51 for i in range(10): #在此处可以选择需要验证的图片，这里博主选择了前10张

52     x,y = test_dataset[i][0],test_dataset[i][1]

53     #可视化

54     show(x).show()

55     # 扩展张量维度为4维

56     x = Variable(torch.unsqueeze(x,dim=0).float(),requires_grad=False).to(device)

57     with torch.no_grad():

58         pred = model(x)

59         # 得到预测类别中最高的那一类，再把最高的这一类对应的标签输出

60         predicted,actual = classes[torch.argmax(pred[0])],classes[y]

61         print(f'predicted:"{predicted},actual:{actual}"')

测试结果，可以看到还是非常不错的

到这手写识别算法基本就完成了

总结

手写识别算，利用现在的框架复现还是比较容易的，代码也是容易读懂，希望这篇博客对你有用

最后的最后，码字不易，给个赞吧wuwuwu~

巴特西

PyTorch复现LeNet-5手写识别学习笔记