LeNet

model.py

import torch.nn as nn
import  torch.nn.functional as F

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet,self).__init__()
        self.conv1=nn.Conv2d(3,16,5)
        self.pool1=nn.MaxPool2d(2,2)
        self.conv2=nn.Conv2d(16,32,5)
        self.pool2=nn.MaxPool2d(2,2)
        self.fc1=nn.Linear(32*5*5,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)
        x=x.view(-1,32*5*5)
        x=F.relu(self.fc1(x))
        x=F.relu(self.fc2(x))
        x=self.fc3(x)
        return x

train.py

import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

transform=transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))]
)
# 50000张训练图片
trainset=torchvision.datasets.CIFAR10(root='./data',train=True,download=False,transform=transform)

trainloader=torch.utils.data.DataLoader(trainset,batch_size=36,shuffle=True,num_workers=0)

# 10000张测试图片
testset=torchvision.datasets.CIFAR10(root='./data',train=False,download=True,transform=transform)

testLoader=torch.utils.data.DataLoader(testset,batch_size=10000,shuffle=True,num_workers=0)

test_data_iter=iter(testLoader)
test_image,test_label=test_data_iter.next()

classes=('plane','car','bird','cat','deer','dog','frog','horse','ship','truck')

# def imshow(img):
#     img = img / 2 + 0.5     # unnormalize  反标准化
#     npimg = img.numpy()
#     plt.imshow(np.transpose(npimg, (1, 2, 0)))
#     plt.show()
#
# # print labels
# print(' '.join(f'{classes[test_label[j]]:5s}' for j in range(4)))
# # show images
# imshow(torchvision.utils.make_grid(test_image))
#


net = LeNet()
loss_function=nn.CrossEntropyLoss()
optimizer=optim.Adam(net.parameters(),lr=0.001)

for epoch in range(5):
    running_loss=0.0
    for step,data in enumerate(trainloader,start=0):
        inputs,labels=data
        outputs=net(inputs)
        loss=loss_function(outputs,labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        running_loss+=loss.item()
        if step%500==499:
            with torch.no_grad():
                outputs=net(test_image)
                predict_y=torch.max(outputs,dim=1)[1]
                accuracy=(predict_y==test_label).sum().item()/test_label.size(0)
                print('[%d, %5d train_loss: %.3f test_accuracy: %.3f' % (epoch+1,step+1,running_loss/500,accuracy))
                running_loss=0.0
print('Finished Training')

save_path='./LeNet.pth'
torch.save(net.state_dict(),save_path)

predict.py

import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet

transform=transforms.Compose(
    [transforms.Resize((32,32)),
     transforms.ToTensor(),
     transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))]
)

classes=('plane','car','bird','cat','deer','dog','frog','horse','ship','truck')

net = LeNet()
net.load_state_dict(torch.load('LeNet.pth'))

im = Image.open('1.jpeg')
im=transform(im)
im=torch.unsqueeze(im,dim=0)  # 增加Batch维度

with torch.no_grad():
    outputs=net(im)
    # predict=torch.max(outputs,dim=1)[1].data.numpy()
    # predict=torch.max(outputs,dim=1)[1].item()
    predict=torch.argmax(torch.softmax(outputs,dim=1),dim=1).item()

print(classes[int(predict)])