CNN入门学习笔记（1）LeNet

• model.py文件

from torch import nn

class LeNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.cov1 = nn.Conv2d(3,6,5,1,0)
        self.pool1 = nn.MaxPool2d(2,2,0)
        self.cov2 = nn.Conv2d(6,16,5,1,0)
        self.pool2 = nn.MaxPool2d(2,2,0)
        self.relu1 = nn.ReLU(inplace=True)
        self.relu2 = nn.ReLU(inplace=True)
        self.relu3 = nn.ReLU(inplace=True)
        self.relu4 = nn.ReLU(inplace=True)
        self.feture = nn.Sequential(
            self.cov1,
            self.relu1,
            self.pool1,
            self.cov2,
            self.relu2,
            self.pool2
        )

        self.linear1 = nn.Linear(16*5*5,120)
        self.linear2 = nn.Linear(120,84)
        self.linear3 = nn.Linear(84,10)
        self.classficar = nn.Sequential(
            self.linear1,
            self.relu3,
            self.linear2,
            self.relu4,
            self.linear3
        )

    def forward(self,input):
        out = self.feture(input)
        out = out.view(-1,16*5*5)
        out = self.classficar(out)
        return out

在该文件中，对LeNet进行了调整，加入ReLu以及将原来的AvgPool改为了MaxPool。

• train.py

from torch.utils.tensorboard import SummaryWriter
import model
import torchvision
import torch
import torch.nn as nn

train_set = torchvision.datasets.CIFAR10(root='../data', train=True,transform=torchvision.transforms.ToTensor(),
                                         download=True)
train_loader = torch.utils.data.DataLoader(train_set,batch_size=64,shuffle=False)
test_set = torchvision.datasets.CIFAR10(root='../data', train=False,transform=torchvision.transforms.ToTensor(),
                                        download=True)
test_loader = torch.utils.data.DataLoader(test_set,batch_size=64,shuffle=False)

epoch = 60
train_num = 0
test_num = 0

writer = SummaryWriter("../logs")
net = model.LeNet()
loss_fn = nn.CrossEntropyLoss()
optmizer = torch.optim.SGD(net.parameters(),lr=0.01)

for i in range(epoch):
    print("----------------第{}轮训练开始----------------".format(i+1))
    for data in train_loader:
        pics,targets = data
        outputs = net(pics)
        loss = loss_fn(outputs,targets)
        optmizer.zero_grad()
        loss.backward()
        optmizer.step()
        train_num = train_num + 1
        # if train_num %10000 == 0:
        #     print("----->>已训练{}次".format(train_num))

    total_loss = 0
    accuracy = 0

    with torch.no_grad():
        for data in test_loader:
            pics, targets = data
            outputs = net(pics)
            loss = loss_fn(outputs,targets)
            total_loss += loss.item()
            accuracy += (outputs.argmax(1) == targets).sum()

    test_set_size =len(test_set)
    print("整体测试集上的正确率为{}%".format(100*accuracy/test_set_size))
    writer.add_scalar("acc",accuracy,i)

torch.save(net, "../model/Lenet.pth")
print(">>>模型已保存")

基于cifar10数据集对模型进行训练，并保存为Lennet.pth文件，最后在测试集上的准确率保持在60%左右。

• detect.py

import cv2
import torch
from PIL import Image
from torchvision.transforms import transforms

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

img = Image.open("../detect_pic/car2.jpeg")
transform = transforms.Compose([transforms.Resize((32, 32)),transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
img = transform(img)
net = torch.load("../model/Lenet.pth")

with torch.no_grad():
    outputs = net(img)
    predict = torch.max(outputs, dim=1)[1].data.numpy() # 索引即classed中的类别
print(classes[int(predict)])


with torch.no_grad():
    outputs = net(img)
    predict = torch.softmax(outputs,dim=1) # [batch，channel，height，width],这里因为对batch不需要处理

print(predict)

对输入的单张图片进行预测，预测效果如下：