CIFAR-10 数据集实战 —— 构建 LeNet5 神经网络

CIFAR-10 数据集网站

如果从官网下载数据集很慢，可以使用国内的地址 http://ai-atest.bj.bcebos.com/cifar-10-python.tar.gz

MNIST 数据集为 0~9 的数字，而 CIFAR-10 数据集为 10 类物品识别，包含飞机、车、鸟、猫等。照片大小为 32*32 的彩色图片（三通道）。每个类别大概有 6000 张照片，其中随机筛选出 5000 用来 training，剩下的 1000 用来 testing

首先引入数据集

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

batch_size=32

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

cifar_test = datasets.CIFAR10(root='cifar', train=False, transform=transforms.Compose([
    transforms.Resize([32, 32]),
    transforms.ToTensor(),
]), download=True)
    
cifar_test = DataLoader(cifar_test, batch_size=batch_size, shuffle=True)
    
x, label = iter(cifar_train).next()
print('x:', x.shape, 'label:', label.shape)

引入数据集以后，接下来开始编写经典的 LeNet5 神经网络

import torch
from torch import nn, optim
import torch.nn.functional as F

class LeNet5(nn.Module):
    """
    for CIFAR10 datasets
    """
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv_unit = nn.Sequential(
            # x: [batchsize, 3, 32, 32] => [batchsize, 6, 28, 28]
            nn.Conv2d(in_channels=3, out_channels=6, kernel_size=5, stride=1, padding=0),
            # [batchsize, 6, 28, 28] => [batchsize, 6, 14, 14]
            nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
            # [batchsize, 6, 14, 14] => [batchsize, 16, 10, 10]
            nn.Conv2d(6, 16, 5, 1, 0),
            # [batchsize, 16, 10, 10] => [batchsize, 16, 5, 5]
            nn.AvgPool2d(2, 2, 0)
            
        )
        
        # fc_unit
        self.fc_unit = nn.Sequential(
            nn.Linear(in_features=16*5*5, out_features=120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 10)
        )        
        
    def forward(self, x):
        batchsize = 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(batchsize, -1)
        
        # [b, 16*5*5] => [b, 10]
        logits = self.fc_unit(x)
        
        return logits
        
def main():

    ##########  train  ##########
    #device = torch.device('cuda')
    #model = LeNet5().to(device)
    criteon = nn.CrossEntropyLoss()
    model = LeNet5()
    optimizer = optim.Adam(model.parameters(), 1e-3)
    for epoch in range(1000):
        model.train()
        for batchidx, (x, label) in enumerate(cifar_train):
            #x, label = x.to(device), label.to(device)
            logits = model(x)
            # logits: [b, 10]
            # label:  [b]
            loss = criteon(logits, label)
            
            # backward
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        
        print('train:', epoch, loss.item())
        
        ########## test  ##########
        model.eval()
        with torch.no_grad():
            total_correct = 0
            total_num = 0
            for x, label in cifar_test:
                # x, label = x.to(device), label.to(device)

                # [b]
                logits = model(x)
                # [b]
                pred = logits.argmax(dim=1)
                # [b] vs [b]
                total_correct += torch.eq(pred, label).float().sum().item()
                total_num += x.size(0)
            acc = total_correct / total_num
            print('test:', epoch, acc)

if __name__ == '__main__':
    main()

从这一部分的运行情况来看，准确率在慢慢上升，但并不稳定，读者有兴趣可以尝试自己修改网络结构，使其准确率更高