【PyTorch】在CIFAR10数据集上实现卷积神经网络LeNet-5

LeNet-5介绍

1. 定义

   • LeNet 诞生于 1994 年，由 Yann Lecun 提出，是一种经典的卷积神经网络，是现代卷积神经网络的起源之一。Yann将该网络用于邮局的邮政的邮政编码识别，有着良好的学习和识别能力。LeNet 又称 LeNet-5，具有一个输入层，两个卷积层，两个池化层，3个全连接层（其中最后一个全连接层为输出层）

2. 网络结构图

   

3. 原理

   • C1 卷积层由 6 个大小为 5* 5 的不同类型的卷积核组成，卷积核的步长为 1，没有零填充，卷积后得到 6 个 28* 28 像素大小的特征图；
   • S2 为最大池化层，池化区域大小为 2* 2，步长为 2，经过S2 池化后得到6 个 14* 14 像素大小的特征图；
   • C3 卷积层由 16 个大小为 5* 5 的不同卷积核组成，卷积核的步长为 1，没有零填充，卷积后得到 16 个 10* 10 像素大小的特征图；
   • S4 最大池化层，池化区域大小为 2* 2，步长为 2，经过 S2 池化后得到 16 个 5* 5 像素大小的特征图；
   • C5 卷积层由120 个大小为 5* 5 的不同卷积核组成，卷积核的步长为 1，没有零填充，卷积后得到 120 个 1* 1 像素大小的特征图；将120 个 1* 1 像素大小的特征图拼接起来作为 F6 的输入；
   • F6 为一个由 84 个神经元组成的全连接隐藏层，激活函数使用sigmoid 函数；
   • 最后一层输出层是一个由10 个神经元组成的 softmax 高斯连接层，可以用来做分类任务

CIFAR10数据集

1. 简介

   • CIFAR-10 是由 Hinton 的学生 Alex Krizhevsky 和 Ilya Sutskever 整理的一个用于识别普适物体的小型数据集。一共包含 10 个类别的 RGB 彩色图 片：飞机（ airplane ）、汽车（ automobile ）、鸟类（ bird ）、猫（ cat ）、鹿（ deer ）、狗（ dog ）、蛙类（ frog ）、马（ horse ）、船（ ship ）和卡车（ truck ）。图片的尺寸为 32×32 ，数据集中一共有 50000 张训练圄片和 10000 张测试图片

2. 示例

   

代码实现

1. lenet5.py

   import torch
   from torch import nn
   
   class Lenet5(nn.Module):
       # for cifar10 dataset
       def __init__(self):
           super(Lenet5, self).__init__()
   
           # 网络结构
           self.conv_unit = nn.Sequential(
               # 输入x:[b, 3, 32, 32]
               # 不加 padding
               # 第一层，输入为 [b, 6, ]
               nn.Conv2d(3, 6, kernel_size = 5, stride = 1, padding = 0),
               nn.ReLU(),
               # Pooling层
               nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
               # 第二个卷积层
               nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),
               nn.ReLU(),
               nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
           )
   
           # 打平操作 flatten
           self.fc_unit = nn.Sequential(
               nn.Linear(16 * 5 * 5, 120),
               nn.ReLU(),
               nn.Linear(120, 84),
               nn.ReLU(),
               nn.Linear(84, 10)
           )
   
           # 计算Linear的第一个参数
           # 假的batch [2, 3, 32, 32] -> [b, 16, 5, 5]
           tmp = torch.randn(2, 3, 32, 32)
           out = self.conv_unit(tmp)
           # torch.Size([2, 16, 5, 5])
           print('conv out:', out.shape)
   
           # 评价标准 use Cross Entropy Loss，分类问题一般使用Cross Entropy
           # self.criteon = nn.CrossEntropyLoss()
   
   
       def forward(self, x):
           # x:[b, 3, 32, 32]
           batchsz = x.size(0)
           # [b, 3, 32, 32] -> [b, 16, 5, 5]
           x = self.conv_unit(x)
           # [b, 16, 5, 5] -> [b, 16*5*5]
           x = x.view(batchsz, 16*5*5)  # 后面也可以直接写-1，自动推算
           # [b, 16*5*5] -> [b, 10]
           logits = self.fc_unit(x)
           return logits
   
   
   # 测试
   def main():
       net = Lenet5()
       tmp = torch.randn(2, 3, 32, 32)
       out = net(tmp)
       print('lenet out:', out.shape)
   
   
   if __name__ == '__main__':
       main()
   

2. CNN.py

   import torch
   from torch.utils.data import DataLoader
   from torchvision import datasets
   from torch import nn, optim
   from torchvision import transforms
   from lenet5 import Lenet5
   
   
   
   def main():
       batchsz = 32
       # 加载CIFAR10数据集
       cifar_train = datasets.CIFAR10('cifar', True, transform = transforms.Compose([
           transforms.Resize([32, 32]),
           transforms.ToTensor()
       ]), download = True)
   
       cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)
   
       cifar_test = datasets.CIFAR10('cifar', False, transform=transforms.Compose([
           transforms.Resize([32, 32]),
           transforms.ToTensor()
       ]), download=True)
   
       cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True)
   
       # 测试
       x, label = iter(cifar_train).next()
       print('x:', x.shape, 'label:', label.shape)
   
   
       device = torch.device('cuda')
       model = Lenet5().to(device)
       criteon = nn.CrossEntropyLoss().to(device)
       optimizer = optim.Adam(model.parameters(), lr=1e-3)
       print(model)
   
       for epoch in range(100):
           model.train()
           for batchidx, (x, label) in enumerate(cifar_train):
               # print(batchidx)
               x, label = x.to(device), label.to(device)
               # logits:[b, 10]
               # label:[b]
               # loss:tensor scalar
               logits = model(x)
               loss = criteon(logits, label)
   
               # backprop
               optimizer.zero_grad()
               loss.backward()
               optimizer.step()
   
           # 完成一个epoch
           print(epoch, "loss = ", loss.item())
   
           model.eval()
           with torch.no_grad():
               # test
               total_correct = 0
               total_num = 0
               for x, label in cifar_test:
                   x, label = x.to(device), label.to(device)
   
                   logits = model(x)
                   # 得到最大值所在的索引
                   pred = logits.argmax(dim=1)
                   # [b] vs [b] -> scalar tensor
                   total_correct += torch.eq(pred, label).float().sum().item()
                   total_num += x.size(0)
               acc = total_correct / total_num
               print(epoch, acc)
   
   
   if __name__ == '__main__':
       main()
   

结果展示

总共进行了 100 轮试验，平均准确率大概只有 60% 左右，简单的 LeNet 网络对较复杂的图片准确率不高