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()

  • 3
    点赞
  • 28
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

让我看看谁在学习

你的鼓励是我创作动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值