使用LeNet5对CIFAR10数据集进行识别分类（pytorch）

张十八员外

已于 2022-10-24 14:00:57 修改

阅读量665

点赞数 1

分类专栏：深度学习入门文章标签： 1024程序员节

于 2022-10-24 13:56:41 首次发布

本文链接：https://blog.csdn.net/weixin_46321755/article/details/127491273

版权

深度学习入门专栏收录该内容

7 篇文章 0 订阅

订阅专栏

该博客介绍了如何使用PyTorch构建经典的LeNet5卷积神经网络模型，并用CIFAR10数据集进行训练。首先，通过torchvision加载和预处理CIFAR10数据，然后定义LeNet5模型结构，接着设置损失函数和优化器。最后，执行训练和测试过程，展示了训练和测试的损失及准确率。

摘要由CSDN通过智能技术生成

LeNet5是一个最早的深度模型，结构如图：
在这里插入图片描述

CIFAR10 数据集包含 60000 张RGB3通道彩色图片，每个图片大小为32*32，这些图像分为 10 个类别，每个类别 6000 张。训练图片共50000 张，测试图片 10000 张。

四步：1、准备数据集；2、定义模型；3、构造损失和优化器；4、训练和测试

代码如下：

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


# Step 1 : prepare dataset

batch_size = 32
cifar_train = datasets.CIFAR10("cifar", train=True, transform=transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor()
]), download=True)
cifar_train_loader = DataLoader(cifar_train, batch_size=batch_size, shuffle=True,)

cifar_test = datasets.CIFAR10("cifar", train=False, transform=transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor()
]), download=True)
cifar_test_loader = DataLoader(cifar_test, batch_size=batch_size, shuffle=False, )


# Step2: design model
class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv_unit = nn.Sequential(
            # x [b, 3, 32, 32] ----> [b, 16, 5, 5]
            nn.Conv2d(3, 6, kernel_size=5, stride=1, padding=0),
            nn.AvgPool2d(kernel_size=2, stride=2, padding=0),

            nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),
            nn.AvgPool2d(kernel_size=2, stride=2,  padding=0)
        )
        # fc layer
        self.fc_unit = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        """
        完成前向传播
        :param x:  [b, 3, 32, 32]
        :return: x: [b, 10]
        """
        batch_size = x.size(0)  # 取x的第一个维度
        # 实现：[b, 3, 32, 32] ----> [b, 16, 5, 5]
        x = self.conv_unit(x)
        # 实现flatten打平 [b, 16, 5, 5] ----> [b, 16*5*5]
        x = x.view(batch_size, -1)
        # 实现: [b, 16*5*5] -----> [b, 10]
        x = self.fc_unit(x)
        return x

model = LeNet5()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)


# Step3: construct Loss and Optimizer
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)


# Step4: Train and Test
def train(epoch):
    running_loss = 0
    model.train()  # 设置为train模式
    for batch_idx, (x, label) in enumerate(cifar_train_loader, 0):
        x, label = x.to(device), label.to(device)
        optimizer.zero_grad()
        # forward
        outputs = model(x)
        loss = criterion(outputs, label)
        # backward
        loss.backward()
        # update
        optimizer.step()
    print("Epoch: ", epoch, "Loss is: ", loss.item())

def test(epoch):
    correct = 0
    total = 0
    model.eval()  # 设置为test模式
    with torch.no_grad():  # 以下内容不需要构建计算图，不需要计算梯度 这一句可加可不加
        for data in cifar_test_loader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            pred = outputs.argmax(dim=1)
            total += labels.size(0)  # 每次循环都把这一批的batch_size加上，就得到总的数量
            correct += torch.eq(pred, labels).float().sum().item()  # 对比预测和label相同的数量 即为预测正确的数量
    print("Epoch", epoch, "Accuracy on test set: %d %%" % (100 * correct / total))
    return correct / total


if __name__ == "__main__":

    for epoch in range(50):
        train(epoch)
        acc = test(epoch)