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

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)