CV经典模型——LeNet-5

本篇文章主要以代码的形式展现，需要注意的地方皆在代码中以注释的方式进行说明，如有疑问欢迎在评论区留言，大家一起进步！
LeNet-5是上世纪末计算机视觉领域的重要产物，其诞生为后来形形色色CNN的产生提供了可能，因此理解其结构精髓是十分重要的，若能自己动手进行实现则更有助于我们的理解，首先上一张LeNet-5的结构图：

结构分析：

从上图中可以很清晰的看出，LeNet-5是一个“2+3”结构（忽略Subsampling层），由两层卷积与三层全连接构成。
注意当时的LeNet使用的sampling方法只是简单的隔行采样，而不是现今使用的更多的MaxPooling或者AvePooling（在后面的代码实现部分使用的是MaxPooling技术）

代码实现（基于pytorch）：

为了尽可能的还原LeNet-5的结构，使用CIFAR-10数据集，同时维度信息等都与原文进行了对应设计，相关部分有进行标注

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

# 定义transform，transform用于对图像进行处理，在Data Augmentation中也会经常使用
my_transform = transform.Compose([transform.Resize((32, 32)), transform.ToTensor()])

# 定义batch_size
batch_size = 32

# 数据集引入
train_set = datasets.CIFAR10(root='dataset', train=True, transform=my_transform, download=True)
train_loader = DataLoader(train_set, batch_size, shuffle=True)

test_set = datasets.CIFAR10(root='dataset', train=False, transform=my_transform, download=True)
test_loader = DataLoader(test_set, batch_size, shuffle=True)


# 自定义一个展平模块
# 在LeNet-5中有从卷积网络到全连接网路的过渡，需要对维度进行压缩
class Flatten(nn.Module):
    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, input):
        return input.view(input.size(0), -1)


# 定义LeNet-5
class Lenet5(nn.Module):
    def __init__(self):
        super().__init__()
        self.module = nn.Sequential(
            # 第一卷积层，[b, 3, 32, 32] => [b, 6, 28, 28]
            nn.Conv2d(3, 6, 5, 1),
            # 第一pool层，这里没有采用Lenet5中的sub_sampling,直接使用max_pool
            # [b, 6, 28, 28] => [b, 6, 14, 14]
            nn.MaxPool2d(2, 2),
            # 第二卷积层,[b,6,14,14] => [b,16,10,10]
            nn.Conv2d(6, 16, 5, 1),
            # 第二pool层，[b,16,10,10] => [b, 16, 5, 5]
            nn.MaxPool2d(2, 2),
            # 展平层,[b, 16*5*5]
            Flatten(),
            # 第一线性层, [b, 16*5*5] => [b, 120]
            nn.Linear(16 * 5 * 5, 120),
            nn.ReLU(),
            # 第二线性层, [b, 120] => [b, 84]
            nn.Linear(120, 84),
            nn.ReLU(),
            # 输出层, [b, 84] => [b, 10]
            nn.Linear(84, 10)
        )

    def forward(self, x):
        logits = self.module(x)
        return logits


def main():
    device = torch.device("cuda")
    # 定义结构对象
    net = Lenet5()
    net.to(device)
    # 定义损失函数
    criterion = nn.CrossEntropyLoss()
    # 定义优化器
    optimizer = optim.Adam(net.parameters(), lr=1e-3)

    # 打印网络结构
    print(net)

    # 进行训练
    for epoch in range(20):
        net.train()
        for batch_index, (data, labels) in enumerate(train_loader):
            # 将数据转移到GPU中进行加速
            data, labels = data.to(device), labels.to(device)

            # 得到预测结果, [b, 10]
            logits = net.forward(data)
            # 误差计算, logits: [b, 10], labels: [b]
            loss = criterion(logits, labels)

            # 计算梯度
            # 梯度清零
            optimizer.zero_grad()
            # 梯度计算
            loss.backward()
            # 传播
            optimizer.step()

        # 输出每次epoch之后的结果
        print("epoch:", epoch, ", loss:", loss.item())

        # 进行测试
        net.eval()
        with torch.no_grad():
            total_num = 0
            correct = 0
            for inputs, labels in test_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                logits = net(inputs)
                pred = logits.argmax(dim=1)
                correct += torch.eq(pred, labels).float().sum().item()
                total_num += inputs.size(0)

            acc = correct / total_num
            print("acc:", acc)


if __name__ == '__main__':
    main()

以上即为本篇文章的全部内容！
转载请注明出处，感谢！