LeNet实现手写数字识别

MNIST数据集：https://github.com/wwjjss/MNIST/tree/main

代码：https://github.com/wwjjss/LeNet_MNIST/tree/main

LeNet介绍：

网络结构：

1：输入，大小表示为（C*H*W）=（channel，hight，width），下同

输入图像大小为32*32

2：C1卷积层

输入通道1，输入通道6，卷积核大小5*5，步长1，填充2，得到输出特征图大小（6*28*28）

3：S2池化层

输入大小6*28*28，卷积核2*2，步长2，使得特征图变为原来的一半，输出大小为（6*14*14）

4：C3卷积层

输入通道6，输出通道16，输入大小（6*14*14），卷积核5*5，输出大小（16*10*10）

5：S4池化层

输入大小16*10*10

代码部分：

import torch
import torch.nn as nn  # 引入pytorch函数库
import torchvision  # 引入pytorch的图形库（torchvision.datasets/.models/.transforms/.utils）
import torchvision.transforms as transforms  # 图形变换（如tensor转换）

# 是否使用GPU
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 设置超参数
NUM_CLASSES = 10  # 分类数目
EPOCHS = 5  # 训练轮次
BATCH_SIZE = 100  # 一轮训练批量大小
LR = 0.001  # 学习率learning_rate

# 创建数据集(root：设置数据集保存路径，train:是否为训练数据集，transforms:将PIL或numpy.ndarray转换为tensor格式,download:是否下载数据集)
train_dataset = torchvision.datasets.MNIST(root='data/', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = torchvision.datasets.MNIST(root='data/', train=False, transform=transforms.ToTensor(), download=False)

# 加载数据集(torch.utilsdata.DataLoader:torch的数据集加载器,dataset:要加载的数据集,batch_size:批量大小,shuffle:是否打乱)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, shuffle=False)


# 构造网络模型
# input_size(1*28*28)
# 卷积：(n-k+2p)/s+1   池化：(n-k)/s+1
class LeNet(nn.Module):
    def __init__(self, num_class=10):
        super(LeNet, self).__init__()  # 继承父类所有属性和方法，父类属性用父类的方法初始化
        self.conv1 = nn.Sequential(  # nn.Sequential:相当于一个容器，将一系列操作包含其中
            nn.Conv2d(1, 6, 5, 1, 2),  # Conv2d(in_channels,out_channels,kernel_size,stride,padding)  out_size(6*28*28)
            nn.ReLU(),  # ReLu()激活函数引入非线性，把负值变为0，正值不变
            nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层，MaxPool2d(kernel_size,stride,padding)  out_size(6*14*14)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(6, 16, 5),  # out_size(16*10*10)
            nn.ReLU(),
            nn.MaxPool2d(2, 2)  # out_size(16*5*5)
        )
        self.fc1 = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),  # 全连接层nn.Linear(in_features,out_features)输入和输出的二维张量大小或神经元个数
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(120, 84),
            nn.ReLU()
        )
        self.fc3 = nn.Linear(84, 10)  # 最后一层得到要数字分类的10类概率值

    # 前向传播

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size()[0], -1)  # 表示将多维的tensor数据展平成一维，torch.view(a,b):重构成a*b维的张量 torch.view(a,-1):-1表示列需要自动计算列数
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x


model = LeNet(NUM_CLASSES).to(DEVICE)  # 将模型加载到设备上

# 构造损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数
optimizer = torch.optim.Adam(model.parameters(), lr=LR)

# 训练模型
total_step = len(train_loader)  # 总训练次数
for epoch in range(EPOCHS):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(DEVICE)
        labels = labels.to(DEVICE)

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播,更新优化器
        optimizer.zero_grad()  # 梯度置零
        loss.backward()  # loss反向传播计算梯度
        optimizer.step()  # 更新网络参数

        if (i + 1) % 100 == 0:
            print(f'Epoch[{epoch + 1}/{EPOCHS}], Step[{i + 1}/{total_step}], Loss:{loss.item():.4f}')

# 测试模型
model.eval()  # 切换到评估模式而非训练模式即固定参数
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(DEVICE)
        labels = labels.to(DEVICE)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)  # 下划线没有实际意义，表示的是具体的value,用其他变量表示也可以，1表示输出所在行的最大值，为0时表示输出所在列的最大值
        total += labels.size(0)  # 数据总量
        correct += (predicted == labels).sum().item()  # 总准确个数
    print(f'Accuracy:{(100 * correct / total)}%')

# 模型保存
torch.save(model.state_dict(), 'LeNet.pth')

测试图片代码：

import torch
import torch.nn as nn
import torchvision
import matplotlib.pyplot as plt
from PIL import Image
from torchvision import transforms
import numpy as np


class LeNet(nn.Module):
    def __init__(self, num_class=10):
        super(LeNet, self).__init__()  # 继承父类所有属性和方法，父类属性用父类的方法初始化
        self.conv1 = nn.Sequential(  # nn.Sequential:相当于一个容器，将一系列操作包含其中
            nn.Conv2d(1, 6, 5, 1, 2),  # Conv2d(in_channels,out_channels,kernel_size,stride,padding)  out_size(6*28*28)
            nn.ReLU(),  # ReLu()激活函数引入非线性，把负值变为0，正值不变
            nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层，MaxPool2d(kernel_size,stride,padding)  out_size(6*14*14)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(6, 16, 5),  # out_size(16*10*10)
            nn.ReLU(),
            nn.MaxPool2d(2, 2)  # out_size(16*5*5)
        )
        self.fc1 = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),  # 全连接层nn.Linear(in_features,out_features)输入和输出的二维张量大小或神经元个数
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(120, 84),
            nn.ReLU()
        )
        self.fc3 = nn.Linear(84, 10)  # 最后一层得到要数字分类的10类概率值

    # 前向传播

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size()[0], -1)  # 表示将多维的tensor数据展平成一维，torch.view(a,b):重构成a*b维的张量 torch.view(a,-1):-1表示列需要自动计算列数
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x


# 读取图片
test_img = 'newTest6.png'
img = plt.imread(test_img)
images = Image.open(test_img)
images = images.resize((28, 28))
images = images.convert('L')
transform = transforms.ToTensor()
images = transform(images)
images = images.resize(1, 1, 28, 28)

# 加载网络模型
model = LeNet()
model.load_state_dict(torch.load('LeNet.pth'))
# model = torch.load('LeNet.pth')
model.eval()
outputs = model(images)

values, indices = outputs.data.max(1)  # 返回最大概率值和下标
plt.title('{}'.format((int(indices[0]))))
plt.imshow(img)
plt.show()