使用LeNet网络进行CIFAR10分类

使用LeNet网络进行多分类问题的实战

这里对cifar10数据集进行了切分，将训练集划分为训练集和验证集。定义了predict函数，可以使用自己下载的图片送入训练好的网络进行识别。

val_set,train_set = random_split(trainset,[40000,10000])

 random_split()是pytorch内置的进行数据集划分函数

import torch
from PIL import Image
from torch.utils.data import random_split
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import torch.nn.functional as F
import torch.optim as optim

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3,out_channels=16,kernel_size=5,stride=1)
        self.pool1 = nn.MaxPool2d(kernel_size=2,stride=2)
        self.conv2 = nn.Conv2d(in_channels=16,out_channels=32,kernel_size=3,stride=1)
        self.pool2 = nn.MaxPool2d(kernel_size=2,stride=2)
        self.fc1 = nn.Linear(in_features=6*6*32,out_features=120)
        self.fc2 = nn.Linear(120,84)
        self.fc3 = nn.Linear(84,10)

    def forward(self,x):
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        #print(x.shape)
        x = x.view(-1,32*6*6)
        #print(x.shape)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x


transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
trainset = torchvision.datasets.CIFAR10(root='./data',train=True,
                                        download=False,transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data',train=False,
                                        download=False,transform=transform)
val_set,train_set = random_split(trainset,[40000,10000])#划分出验证集（10000），和训练集（40000）
#print(len(val_set))

#使用验证集进行训练
trainloader = torch.utils.data.DataLoader(dataset=val_set,batch_size=32,shuffle=True,num_workers=0,drop_last=True)
testloader = torch.utils.data.DataLoader(dataset=testset,batch_size=5000,num_workers=0)
val_data_iter = iter(testloader)
val_image, val_label = val_data_iter.next()
classes = ('plan','car','bird','cat','deer','dog','frog','horse','ship','truck')
model = LeNet()
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(),lr=0.001)

def train():
    for epoch in range(5):
        run_loss = 0.0
        for step,data in enumerate(trainloader):
            inputs,labels = data
            #print(inputs.size())
            optimizer.zero_grad()
            output = model(inputs)
            loss = loss_function(output,labels)
            loss.backward()
            optimizer.step()
            run_loss += loss.item()
            if step % 500 == 499:
                with torch.no_grad():
                    output = model(val_image)
                    predict_y = torch.argmax(output,dim=1)
                    #print(predict_y)
                    accuracy = torch.eq(predict_y,val_label).sum().item()/val_label.size(0)
                    print('[%d, %5d] train_loss: %.3f  test_accuracy: %.3f' %
                          (epoch + 1, step + 1, run_loss / 500, accuracy))
                    run_loss = 0.0
    print('finished train')
    save_path = './LeNet.pth'
    torch.save(model.state_dict(),save_path)
def predict():
    transform = transforms.Compose(
        [transforms.Resize((32, 32)),
         transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    net = LeNet()
    net.load_state_dict(torch.load('Lenet.pth'))

    im = Image.open('cat.jpg')
    im = transform(im)  # [C, H, W]
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]

    with torch.no_grad():
        outputs = net(im)
        predict = torch.max(outputs, dim=1)[1].data.numpy()
    print(classes[int(predict)])
if __name__=='__main__':
    #train()
    predict()