欢迎访问宙启技术站
智能推送

使用Python编写LeNet()模型进行手写数字识别

发布时间:2023-12-24 16:26:34

LeNet是一个经典的卷积神经网络模型,最早由Yann LeCun在1998年设计用于手写数字识别。该模型在MNIST数据集上取得了很好的效果,并且奠定了卷积神经网络的基本架构。

下面我将介绍如何使用Python编写LeNet()模型进行手写数字识别,并给出一个简单的使用例子。

首先,我们需要导入必要的库:

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

然后,我们可以定义LeNet()模型的结构:

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.fc1 = nn.Linear(16*4*4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = x.view(-1, 16*4*4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

LeNet()模型的结构包括两个卷积层和三个全连接层。其中,卷积层(conv)用于提取图像的特征,全连接层(fc)用于分类。

下一步是定义训练和测试的函数:

def train(model, device, train_loader, optimizer, criterion):
    model.train()
    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()

def test(model, device, test_loader):
    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 += F.cross_entropy(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    
    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    return test_loss, accuracy

训练函数(train)通过循环遍历训练数据集进行模型的训练,并更新模型的参数。测试函数(test)则通过遍历测试数据集计算模型在测试集上的损失值和准确率。

最后,我们可以创建一个训练和测试的主函数来调用上述的函数:

def main():
    # 设置随机种子
    torch.manual_seed(0)
    
    # 定义设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 加载数据集
    train_loader = torch.utils.data.DataLoader(
        torchvision.datasets.MNIST('data', train=True, download=True,
                                   transform=torchvision.transforms.Compose([
                                       torchvision.transforms.ToTensor(),
                                       torchvision.transforms.Normalize((0.1307,), (0.3081,))
                                   ])),
        batch_size=64, shuffle=True)

    test_loader = torch.utils.data.DataLoader(
        torchvision.datasets.MNIST('data', train=False, transform=torchvision.transforms.Compose([
                                       torchvision.transforms.ToTensor(),
                                       torchvision.transforms.Normalize((0.1307,), (0.3081,))
                                   ])),
        batch_size=1000, shuffle=True)
    
    # 初始化模型
    model = LeNet().to(device)
    
    # 定义优化器和损失函数
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    
    # 训练模型
    for epoch in range(1, 11):
        train(model, device, train_loader, optimizer, criterion)
        test_loss, accuracy = test(model, device, test_loader)
        print('Epoch {}: Test Loss = {:.4f}, Accuracy = {:.2f}%'.format(
            epoch, test_loss, accuracy))

在主函数中,我们首先设置随机种子,然后加载MNIST数据集。接着,初始化LeNet()模型,并定义Adam优化器和交叉熵损失函数。最后,通过循环训练模型,并输出每个epoch的测试损失值及准确率。

现在我们可以调用main()函数开始训练和测试LeNet()模型了:

if __name__ == '__main__':
    main()

这样,我们就完成了使用Python编写LeNet()模型进行手写数字识别的过程,并给出了一个简单的使用例子。该例子可以帮助您更好地理解LeNet模型的结构和训练流程。希望对您有所帮助!