CNN入门学习笔记（3）VGG16

• model.py文件

from torch import nn
from torchvision import models

class Vgg16(nn.Module):
    def __init__(self):
        super().__init__()
        net = models.vgg16(True)
        net.classifier = nn.Sequential()
        self.futures = net
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(25088, 512),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(512, 128),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(128, 10),
        )

    def forward(self,x):
        x=self.futures(x)
        x=x.view(-1,25088)
        x=self.classifier(x)
        return x

这里直接调用了torchvirsion中的vgg网络，并重新编写了全连接层。

• train.py

from torch.utils.tensorboard import SummaryWriter
import model
import torchvision
import torch
import torch.nn as nn
import  torchvision.transforms as transforms
from torchvision.transforms.functional import InterpolationMode
transform = transforms.Compose([transforms.Resize((224, 224), interpolation=InterpolationMode.BICUBIC),
                                transforms.ToTensor(),
                                transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.247, 0.2435, 0.2616]) # 此为训练集上的均值与方差
                                ])
train_set = torchvision.datasets.CIFAR10(root='../data', train=True,transform=transform,
                                         download=True)
train_loader = torch.utils.data.DataLoader(train_set,batch_size=256,shuffle=True)
test_set = torchvision.datasets.CIFAR10(root='../data', train=False,transform=transform,
                                        download=True)
test_loader = torch.utils.data.DataLoader(test_set,batch_size=256,shuffle=True)

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

epoch = 10
train_num = 0
test_num = 0

writer = SummaryWriter("../logs")
net = model.Vgg16().to(device)
# net = torch.load("../model/vgg16.pth",map_location=torch.device('cpu'))
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
        pics = pics.to(device)
        targets = targets.to(device)
        outputs = net(pics)
        loss = loss_fn(outputs,targets)
        optmizer.zero_grad()
        loss.backward()
        optmizer.step()
        train_num = train_num + 1
        if train_num %100 == 0:
            print("----->>已训练{}次".format(train_num))

    total_loss = 0
    accuracy = 0

    with torch.no_grad():
        for data in test_loader:
            pics, targets = data
            pics = pics.to(device)
            targets = targets.to(device)
            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/vgg16.pth")
print(">>>模型已保存")

首先将cifar10数据集通过transform转化为224x224大小，并进行toTensor变换，再对模型进行训练。训练部分结果如下：

在训练集上的准确度已经达到了90%以上，超过LeNet和AlexNet。

• detect.py

import cv2
import torch
from PIL import Image
from torchvision.transforms import transforms
from torchvision.transforms.functional import InterpolationMode
import numpy as np
classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')
img = Image.open("../detect_pic/dear1.jpeg")

transform = transforms.Compose([transforms.Resize((224, 224), interpolation=InterpolationMode.BICUBIC),
                                transforms.ToTensor(),
                                transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.247, 0.2435, 0.2616]) # 此为训练集上的均值与方差
                                ])
img = transform(img)
net = torch.load("../model/vgg16.pth",map_location = 'cpu')

# net.eval()
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)

im = cv2.imread("../detect_pic/dear1.jpeg")
cv2.imshow("",im)
cv2.waitKey()

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