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

Pytorch实现LeNet-5:深度学习中的经典之作

发布时间:2023-12-17 08:08:14

深度学习中的经典之作之一就是LeNet-5(LeNet)模型。这是由Yann LeCun于1998年提出的,是最早的卷积神经网络(Convolutional Neural Network, CNN)之一,也被广泛应用于手写数字识别任务。

LeNet-5主要由两个卷积层和三个全连接层组成,其中卷积层用来提取特征,全连接层用来进行分类。在当时的计算机和数据集规模限制下,LeNet模型非常适合处理手写数字识别任务。即使在现在的计算机和数据集规模较大的情况下,LeNet仍然被视为了解卷积神经网络的重要里程碑。

下面我们用PyTorch实现LeNet-5,并使用MNIST手写数字数据集进行训练和测试。

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

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader

接下来,我们定义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 = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, kernel_size=2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, kernel_size=2)
        x = x.view(x.size(0), -1)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

LeNet类继承自nn.Module,并在构造函数中初始化所有卷积层和全连接层。在前向传播方法forward中,我们定义了relu和max_pool2d操作,将输入数据通过各个层,并返回最后一层分类的输出。

现在我们可以加载MNIST数据集并进行预处理:

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = MNIST(root='data', train=True, transform=transform, download=True)
test_dataset = MNIST(root='data', train=False, transform=transform)

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

我们使用了torchvision.transforms中的Compose函数来定义预处理操作,将图像转换为Tensor格式并进行归一化。

接下来,我们可以定义优化器和损失函数,并开始训练模型:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = LeNet().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
criterion = nn.CrossEntropyLoss()

num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print('Epoch {} - Training loss: {:.4f}'.format(epoch+1, running_loss / len(train_loader)))

    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print('Epoch {} - Testing accuracy: {:.2f}%'.format(epoch+1, 100 * correct / total))

在每个训练周期中,我们将模型设置为训练模式(model.train())并进行前向传播、反向传播和模型参数优化。然后,我们评估模型在测试集上的准确率,并打印出训练损失和测试准确率。

通过上述步骤,我们就能够使用PyTorch实现LeNet-5模型,并在MNIST数据集上进行训练和测试。LeNet-5模型是深度学习中的经典之作,对于初学者来说是一个很好的入门模型。