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

用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

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

import torch
from torch import nn

#定义一个网络模型类
class MyLeNet5(nn.Module):
    #初始化网络
    def __init__(self):
        super(MyLeNet5,self).__init__()
        #输入大小为32*32，输出大小为28*28,输入通道为1，输出为6，卷积核大小为5,步长为1
        self.c1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, padding=2)
        #sigmoid激活函数
        self.Sigmoid= nn.Sigmoid()
        #平均池化
        self.s2 = nn.AvgPool2d(kernel_size=2, stride=2)
        self.c3 = nn.Conv2d(in_channels=6,out_channels=16,kernel_size=5)
        self.s4 = nn.AvgPool2d(kernel_size=2,stride=2)
        self.c5 = nn.Conv2d(in_channels=16,out_channels=120,kernel_size=5)
        #展开
        self.flatten = nn.Flatten()
        self.f6 = nn.Linear(120,84)
        self.output = nn.Linear(84,10)

    def forward(self,x):
        #输入x为32*32*1,输出为28*28*6
        x = self.Sigmoid(self.c1(x))
        #输入为28*28*6,输出为14*14*6
        x = self.s2(x)
        # 输入为14*14*6,输出为10*10*16
        x = self.Sigmoid(self.c3(x))
        # 输入为10*10*16,输出为5*5*16
        x = self.s4(x)
        # 输入为5*5*16,输出为1*1*120
        x = self.c5(x)
        x = self.flatten(x)
        # 输入为120，输出为84
        x = self.f6(x)
        # 输入为84，输出为10
        x = self.output(x)
        return x

if __name__=="__main__":
    x = torch.rand([1,1,28,28])#任意产生一个张量,批次1，通道为1，大小为28*28
    model = MyLeNet5()#网络实例化
    y = model(x) #输出结果

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

3.train.py

这是用于训练模型的代码

import torch
from torch import nn
from net import MyLeNet5
from torch.optim import  lr_scheduler
from torchvision import datasets,transforms
import os


#将数据转化为tensor格式
data_transform = transforms.Compose([
    transforms.ToTensor()
])

# 加载训练数据集
train_dataset = datasets.MNIST(root='./data', train=True, transform=data_transform, download=True)
# 给训练集创建一个数据加载器, shuffle=True用于打乱数据集，每次都会以不同的顺序返回。
train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)
# 加载训练数据集
test_dataset = datasets.MNIST(root='./data', train=False, transform=data_transform, download=True)
# 给训练集创建一个数据加载器, shuffle=True用于打乱数据集，每次都会以不同的顺序返回。
test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=16, shuffle=True)


# 如果显卡可用，则用显卡进行训练
device = "cuda" if torch.cuda.is_available() else 'cpu'

#调用net文件的模型，果GPU可用则将模型转到GPU
model = MyLeNet5().to(device)

#定义损失函数，交叉熵损失
loss_fn = nn.CrossEntropyLoss()

#定义优化器SGD,随机梯度下降
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)

#学习率每10个epoch变为原来的0.1
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

#定义训练函数
def train(dataloader, model, loss_fn, optimizer):
    loss, current, n = 0.0, 0.0, 0
    # enumerate返回为数据和标签还有批次
    for batch, (X, y) in enumerate(dataloader):
        # 前向传播
        X, y = X.to(device), y.to(device)
        output = model(X)
        cur_loss = loss_fn(output, y)
        # torch.max返回每行最大的概率和最大概率的索引,由于批次是16，所以返回16个概率和索引
        _, pred = torch.max(output, axis=1)

        # 计算每批次的准确率， output.shape[0]为该批次的多少
        cur_acc = torch.sum(y == pred) / output.shape[0]
        # print(cur_acc)
        # 反向传播
        optimizer.zero_grad()
        cur_loss.backward()
        optimizer.step()
        # 取出loss值和精度值
        loss += cur_loss.item()
        current += cur_acc.item()
        n = n + 1

    print('train_loss:' + str(loss / n))
    print('train_acc:' + str(current / n))


#定义验证函数
def val(dataloader,model,loss_fn):
    # 将模型转为验证模式
    model.eval()
    loss, acc, n = 0.0, 0.0, 0
    # enumerate返回为数据和标签还有批次
    with torch.no_grad():
        for batch, (x, y) in enumerate(dataloader):
            # 前向传播
            x, y = x.to(device), y.to(device)
            output = model(x)
            cur_loss = loss_fn(output, y)
            # torch.max返回每行最大的概率和最大概率的索引,由于批次是16，所以返回16个概率和索引
            _, pred = torch.max(output, axis=1)

            # 计算每批次的准确率， output.shape[0]为该批次的多少
            cur_acc = torch.sum(y == pred) / output.shape[0]
            loss += cur_loss.item()
            acc += cur_acc.item()#取出单元素张量的元素值并返回该值
            n += 1  # 记录有多少批次
        print('test_loss:' + str(loss / n))
        print('test_acc:' + str(acc / n))

        return acc/n

#开始训练
epoch = 30#训练轮次
max_acc = 0
for t in range(epoch):
    lr_scheduler.step()#学习率调整
    print(f"epoch{t+1}\n-------------------")#加f表示格式化字符串，加f后可以在字符串里面使用用花括号括起来的变量和表达式
    train(train_dataloader, model, loss_fn, optimizer)#调用train函数
    a = val(test_dataloader,model,loss_fn)
    #保存最后的模型权重文件
    if a > max_acc:
        folder = 'save_model'
        if not os.path.exists(folder):
            os.mkdir('save_model')
        max_acc = a
        print('save best model')
        torch.save(model.state_dict(),"save_model/best_model.pth")
    #保存最后的文件
    if t == epoch - 1:
        torch.save(model.state_dict(),"save_model/last_model.pth")
print('Done')

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

 

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

三、模型测试

1.test.py

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

import torch
from net import MyLeNet5
from torch.autograd import Variable
from torchvision import datasets,transforms
from torchvision.transforms import ToPILImage

# 将数据转化为tensor格式
data_transform = transforms.Compose([
    transforms.ToTensor()
])

# 加载训练数据集
train_dataset = datasets.MNIST(root='./data', train=True, transform=data_transform, download=True)
# 给训练集创建一个数据加载器, shuffle=True用于打乱数据集，每次都会以不同的顺序返回。
#train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)
# 加载训练数据集
test_dataset = datasets.MNIST(root='./data', train=False, transform=data_transform, download=True)
# 给训练集创建一个数据加载器, shuffle=True用于打乱数据集，每次都会以不同的顺序返回。
#test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=16, shuffle=True)

#  如果显卡可用，则用显卡进行训练
device = "cuda" if torch.cuda.is_available() else 'cpu'

# 调用net里面定义的模型，如果GPU可用则将模型转到GPU
model = MyLeNet5().to(device)

#加载train.py里训练好的模型
model.load_state_dict(torch.load(("D:/python/LeNet-5/save_model/best_model.pth")))#填写权重路径

#获取预测结果

classes = [
    "0",
    "1",
    "2",
    "3",
    "4",
    "5",
    "6",
    "7",
    "8",
    "9",
]

# 把tensor转换成Image，方便可视化
show = ToPILImage()

#进入验证阶段
model.eval()
# 对test_dataset手写数字图片进行推理
for i in range(10): #在此处可以选择需要验证的图片，这里博主选择了前10张
    x,y = test_dataset[i][0],test_dataset[i][1]
    #可视化
    show(x).show()
    # 扩展张量维度为4维
    x = Variable(torch.unsqueeze(x,dim=0).float(),requires_grad=False).to(device)
    with torch.no_grad():
        pred = model(x)
        # 得到预测类别中最高的那一类，再把最高的这一类对应的标签输出
        predicted,actual = classes[torch.argmax(pred[0])],classes[y]
        print(f'predicted:"{predicted},actual:{actual}"')

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

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

总结

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

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