Pytorch用LeNet在CIFAR10数据集上实现分类

本文主要讲解该算法的实现过程

实验环境

python3.6+pytorch1.2+cuda10.1

数据集

CIFAR10数据集包含6万训练图像与1万验证图像

下边是代码实现过程及讲解

数据加载

#载入数据
def loadCIFAR10(batch_size):
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    trainset = CIFAR10('./data',train=True,transform=transform,download=True)
    testset = CIFAR10('./data',train=False,transform=transform,download=True)

    trainloader = DataLoader(trainset,batch_size=batch_size,shuffle=True,num_workers=2)
    testloader = DataLoader(testset,batch_size=batch_size,shuffle=False,num_workers=2)

    return trainset,testset,trainloader,testloader

定义网络结构

class Lenet(nn.Module):
    def __init__(self):
        super(Lenet,self).__init__()
        self.conv1 = nn.Conv2d(3,6,5)
        self.pool = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(6,16,5)
        self.fc1 = nn.Linear(400,120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self,x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        # print(2,x.size())
        x = x.view(-1,400)
        x = F.relu(self.fc1(x))
        # print(3,x.size())
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

参数设定

learning_rate = 1e-3
batch_size = 10
epoches = 50

lenet = Lenet()
lenet.cuda()

trainset,testset,trainloader,testloader = loadCIFAR10(batch_size)
print(len(trainset))
criterian = nn.CrossEntropyLoss()
optimizer = optim.SGD(lenet.parameters(),lr=learning_rate,momentum=0.9)

网络训练

for epoch in range(epoches):
    running_loss = 0.0
    for i,data in enumerate(trainloader,0):
        inputs,labels = data
        inputs = inputs.cuda()
        labels = labels.cuda()
        optimizer.zero_grad()
        # print(0,inputs.size())
        outputs = lenet(inputs)

        loss = criterian(outputs,labels)
        loss.backward()
        optimizer.step()

        running_loss +=loss.item()
        if i%2000==1999:
            print('[%d,%5d] loss: %.3f'%(epoch+1,i+1,running_loss/2000))
            running_loss = 0.0

结果检验

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images,labels = data
        images = images.cuda()
        labels = labels.cuda()
        outputs = lenet(images)
        _,predict = torch.max(outputs.data,1)
        total += labels.size(0)
        correct += (predict==labels).sum().item()
print('Acc: %.2f%%'%(100*correct/total))

结果展示

迭代了50次最后结果也才刚过60%，主要原因是模型太简单了，如果要进一步提高正确率要考虑更换模型

完整代码

import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision import transforms
from torch.autograd import Variable
from torch import optim
import  torch.nn as nn
import  torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import pylab

classes = ['airplane','automobile','bird','cat','deer','dog','frog','horse','ship','truck']


#载入数据
def loadCIFAR10(batch_size):
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    trainset = CIFAR10('./data',train=True,transform=transform,download=True)
    testset = CIFAR10('./data',train=False,transform=transform,download=True)

    trainloader = DataLoader(trainset,batch_size=batch_size,shuffle=True,num_workers=2)
    testloader = DataLoader(testset,batch_size=batch_size,shuffle=False,num_workers=2)

    return trainset,testset,trainloader,testloader


class Lenet(nn.Module):
    def __init__(self):
        super(Lenet,self).__init__()
        self.conv1 = nn.Conv2d(3,6,5)
        self.pool = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(6,16,5)
        self.fc1 = nn.Linear(400,120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self,x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        # print(2,x.size())
        x = x.view(-1,400)
        x = F.relu(self.fc1(x))
        # print(3,x.size())
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

learning_rate = 1e-3
batch_size = 10
epoches = 50

lenet = Lenet()
lenet.cuda()

trainset,testset,trainloader,testloader = loadCIFAR10(batch_size)
print(len(trainset))
criterian = nn.CrossEntropyLoss()
optimizer = optim.SGD(lenet.parameters(),lr=learning_rate,momentum=0.9)


for epoch in range(epoches):
    running_loss = 0.0
    for i,data in enumerate(trainloader,0):
        inputs,labels = data
        inputs = inputs.cuda()
        labels = labels.cuda()
        optimizer.zero_grad()
        # print(0,inputs.size())
        outputs = lenet(inputs)

        loss = criterian(outputs,labels)
        loss.backward()
        optimizer.step()

        running_loss +=loss.item()
        if i%2000==1999:
            print('[%d,%5d] loss: %.3f'%(epoch+1,i+1,running_loss/2000))
            running_loss = 0.0

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images,labels = data
        images = images.cuda()
        labels = labels.cuda()
        outputs = lenet(images)
        _,predict = torch.max(outputs.data,1)
        total += labels.size(0)
        correct += (predict==labels).sum().item()
print('Acc: %.2f%%'%(100*correct/total))