Pytorch 实践手写数字识别深度学习网络 LeNet-5

Pytorch 实践手写数字识别深度学习网络 LeNet-5

训练手写体识别任务是一个非常简单的学习任务，理论简单是简单，但是我相信很多人都和我一样，想体验一下一个学习任务的全流程，有原始数据，处理数据，编写网络，训练模型，测试模型，使用模型这个过程。今天我们就由我来带大家体验一下。

认识 LeNet-5

LeNet-5出自论文《Gradient-Based Learning Applied to Document Recognition》， 原本是一种用于手写体字符识别的非常高效的卷积神经网络，包含了深度学习的基本模块：卷 积层，池化层，全连接层。

• INPUT（输入层） ：输入28∗28的图片。
• C1（卷积层）：选取6个5∗5卷积核(不包含偏置)，得到6个特征图，每个特征
• 图的一个边为28−5+1=24。
• S2（池化层）：池化层是一个下采样层，输出12∗12∗6的特征图。
• C3（卷积层）：选取16个大小为5∗5卷积核，得到特征图大小为8∗8∗16。
• S4（池化层）：窗口大小为2∗2，输出4∗4∗16的特征图。
• F5（全连接层）：120个神经元。
• F6（全连接层）：84个神经元。
• OUTPUT（输出层）：10个神经元，10分类问题。

认识数据集

MNIST数据集来自美国国家标准与技术研究所，National Institute of Standards and Technology(NIST),数据集由来自250个不同人手写的数字构 成，其中50%是高中学生，50%来自人口普查局（the Census Bureau）的工 作人员。

训练集：60000，测试集：10000

MNIST数据集可在 http://yann.lecun.com/exdb/mnist/ 获取

大家如果想要的话可以联系我邮箱2837468248@qq.com，也可以直接发给你。

处理数据集

一般我们进行这个实践的时候，因为这个太经典了，很多的框架就直接集成了这个数据集，不用我们自己处理了，直接拿来用就好了。如下：

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

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

其中的 torchvision 的包的 MNIST 数据集已经替我们解析好了一切。下载、解压、加载都做好了，效果如下：

虽然这样我们也能用，但是我们今天是来体验全流程的，我们要自己处理。

下载数据集

可以通过上面的数据集下载网址进行下载

下载好后是 gz 格式，你可以选择用程序对其进行解压，或者直接用解压软件解压。

解压后我是这样存放数据的

读取数据

一开始看到这个 ubyte 形式的数据我直接震惊，这啥数据呀，哪里有图片呢，真不清楚，后面了解到这是把数据进行压缩了。 具体解释请看文章

读取图片数据：

# 读取图像文件
def load_mnist_images(file_path):
    with open(file_path, 'rb') as f:
        magic, num_images, rows, cols = struct.unpack('>IIII', f.read(16))
        if magic != 2051:
            raise ValueError(f'Invalid magic number {magic} in file: {file_path}')
        images = np.frombuffer(f.read(), dtype=np.uint8).reshape(num_images, 1, rows, cols)
    return images

读取标签数据

# 读取标签文件
def load_mnist_labels(file_path):
    with open(file_path, 'rb') as f:
        magic, num_labels = struct.unpack('>II', f.read(8))
        if magic != 2049:
            raise ValueError(f'Invalid magic number {magic} in file: {file_path}')
        labels = np.frombuffer(f.read(), dtype=np.uint8)
    return labels

定义Dataset的继承类

原来直接用 torch vision 的数据的话，他会直接包装好，但是我们这里是要全部体验全流程，所以我们自己包装。

为什么要有 Dataset 类呢？

因为在进行深度学习训练的时候都是一批一批进行训练的，需要把数据载入到 Pytorch 提供的 dataloader 中去，方便 pytorch 后面对我们的数据方便进行操作。

我们自己定义一个 Dataset 类的话，一定要实现三个函数

__init__

__len__

__getitem__

这里我们的实现如下：

# 自定义 Dataset 类
class MNISTDataset(Dataset):
    def __init__(self, images, labels, transform=None):
        self.images = images
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label

把数据进行载入

这里载入的时候我们有个小地方需要注意一下，对数据进行一个简单的处理，就是把数据的维度进行放小，原来图片的维度如下：

我们需要的数据维度是 (60000,28,28）的数据，所以要进行降维

train_images = train_images.squeeze(axis=1) # 把第一维度进行减掉

后面测试集的数据的话就是一样的处理。

载入dataloader

transform = transforms.Compose([
    transforms.ToTensor(), #把numpy数据转化为tensor
    transforms.Normalize((0.5,), (0.5,)) #对数据进行归一化处理
])
from torch.utils.data import DataLoader
train_dataset = MNISTDataset(train_images, train_labels,transform=transform)
test_dataset = MNISTDataset(test_images,test_labels,transform=transform)

train_loader = DataLoader(dataset=train_dataset, batch_size=64,shuffle=True)
test_loader = DataLoader(dataset=test_dataset,batch_size=64,shuffle=False)

编写网络

这个网络是非常简单的，直接定义一个类继承 torch.nn.Module进行实现就好了,代码如下:

# 网络定义
import torch
import torch.nn as nn
from torch.utils.data import Dataset
import torch.optim as optim
import torchvision.transforms as transforms


class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=2)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x)) #卷积以后进行激活 
        x = torch.max_pool2d(x, kernel_size=2, stride=2) #最大池化，提取特征
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, kernel_size=2, stride=2)
        x = x.view(-1, 16 * 5 * 5) #把数据进行展平方便全连接层的输入
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

编写训练与测试代码

训练的话就是一般深度学习的流程，直接调用 pytorch 的API进行解决了。

# 训练和测试函数
def train(model, device, train_loader, optimizer, criterion, epoch):
    model.train()
    best_model = model
    min_loss = 1
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        if min_loss > loss.item():
            best_model, best_loss = model, loss.item()
            print("update")
        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')
    print("模型训练结束")
    print("保存最好 loss 模型，loss:",min_loss)
    torch.save(best_model.state_dict(),'best-lenet5.pth')

def test(model, device, test_loader, criterion):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.0f}%)\n')
# 训练和测试模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = LeNet5().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

num_epochs = 10
for epoch in range(1, num_epochs + 1):
    train(model, device, train_loader, optimizer, criterion, epoch)
    test(model, device, test_loader, criterion)

实践结果展示

载入保存的模型 pth 文件，然后手写一个 28 * 28 的数字图片进行处理识别效果如下

代码如下：

model = LeNet5()
model.load_state_dict(torch.load('best-lenet5.pth',map_location=torch.device('cpu')))
# 假设图像路径
image_path = '5.png'  # 替换为你的图像路径
from PIL import Image
# 使用 PIL 库打开图像
image = Image.open(image_path)

# 使用 torchvision.transforms 进行数据转换和归一化
transform = transforms.Compose([
    transforms.Resize((28, 28)),  # 调整大小到 28x28
    transforms.Grayscale(),
    transforms.ToTensor(),        # 转为 Tensor
    transforms.Normalize((0.5,), (0.5,))  # 归一化
])

# 应用转换
image_tensor = transform(image).unsqueeze(0)
predict_output = model(image_tensor)
pred_num = predict_output.argmax(dim=1,keepdim=True)
print(pred_num) # 数据要写满28*28的格子才能预测)

如果需要 pth 文件也可以联系我，不过这个训练很快，可以自己训练玩！

完整代码

# 网络定义
import torch
import torch.nn as nn
from torch.utils.data import Dataset
import torch.optim as optim
import torchvision.transforms as transforms


class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=2)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, kernel_size=2, stride=2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, kernel_size=2, stride=2)
        x = x.view(-1, 16 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x
# 自定义 Dataset 类
class MNISTDataset(Dataset):
    def __init__(self, images, labels, transform=None):
        self.images = images
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label
import numpy as np
import struct


# 读取标签文件
def load_mnist_labels(file_path):
    with open(file_path, 'rb') as f:
        magic, num_labels = struct.unpack('>II', f.read(8))
        if magic != 2049:
            raise ValueError(f'Invalid magic number {magic} in file: {file_path}')
        labels = np.frombuffer(f.read(), dtype=np.uint8)
    return labels


# 读取图像文件
def load_mnist_images(file_path):
    with open(file_path, 'rb') as f:
        magic, num_images, rows, cols = struct.unpack('>IIII', f.read(16))
        if magic != 2051:
            raise ValueError(f'Invalid magic number {magic} in file: {file_path}')
        images = np.frombuffer(f.read(), dtype=np.uint8).reshape(num_images, 1, rows, cols)
    return images

# 获取到标签，图像数据
train_images = load_mnist_images('./data/train/train-images-idx3-ubyte')
train_labels = load_mnist_labels('./data/train/train-labels-idx1-ubyte')
print('train_images.shape', train_images.shape)
print('label.shape', train_labels.shape)
train_images = train_images.squeeze(axis=1)
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
test_images = load_mnist_images('./data/test/t10k-images-idx3-ubyte')
test_labels = load_mnist_labels('./data/test/t10k-labels-idx1-ubyte')
test_images = test_images.squeeze(axis = 1)
print(test_images.shape)
print(test_labels.shape)
from torch.utils.data import DataLoader
train_dataset = MNISTDataset(train_images, train_labels,transform=transform)
test_dataset = MNISTDataset(test_images,test_labels,transform=transform)

train_loader = DataLoader(dataset=train_dataset, batch_size=64,shuffle=True)
test_loader = DataLoader(dataset=test_dataset,batch_size=64,shuffle=False)
# 训练和测试函数
def train(model, device, train_loader, optimizer, criterion, epoch):
    model.train()
    best_model = model
    min_loss = 100000.0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        if min_loss > loss.item():
            best_model, best_loss = model, loss.item()
            print("update")
        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')
    print("模型训练结束")
    print("保存最好 loss 模型，loss:",min_loss)
    torch.save(best_model.state_dict(),'best-lenet5.pth')

def test(model, device, test_loader, criterion):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.0f}%)\n')

# 训练和测试模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = LeNet5().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

num_epochs = 10
for epoch in range(1, num_epochs + 1):
    train(model, device, train_loader, optimizer, criterion, epoch)
    test(model, device, test_loader, criterion)