pytorch VGG11识别cifar10数据集(训练+预测单张输入图

1. 简介

本文将介绍如何使用pytorch来构建一个VGG11模型，并将训练好的模型应用于cifar10数据集。我们将使用pytorch预处理模块来加载cifar10数据集，使用torchvision.models模块提供的VGG11模型来进行训练和预测。我们通过训练和预测单张输入图，来展示VGG11模型的性能。

2. 数据集介绍

CIFAR-10是一个经典的图片分类数据集，它包含10个类别的60000张32x32彩色图片。这些类别包括：airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck。每个类别中都有6000张图片。其中50000张被用作训练数据，剩下的10000张被用作测试/验证数据。

import torch

import torchvision
import torchvision.transforms as transforms
#预处理模块
transform = transforms.Compose(
    [transforms.Resize(256), #将图像的最小边缩放到256
     transforms.RandomCrop(224), #抠取224x224的随机patch
     transforms.RandomHorizontalFlip(), #随机水平翻折
     transforms.ToTensor(),  #将PIL图像转换为Tensor
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) #输入数据的均值和标准差
#trainset and testset
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)

我们可以看到，我们使用了transforms模块，其中包括如下几个步骤：

将图像的最小边缩放到256

抠取224x224的随机patch

随机水平翻折

将PIL图像转换为Tensor

对图像进行标准化

3. 模型介绍

VGG是一个经典的卷积神经网络模型，它在2014年ILSVRC比赛中取得了较好的成绩。VGG包含了很多卷积层和池化层，它的结构相对简单，每个卷积层都是一个3x3的卷积，每个池化层都是2x2的最大池化。VGG模型共有16层或者19层，其中11层和13层比较常用，我们将使用VGG11模型。

import torch.nn as nn

import torch.nn.functional as F
#VGG11网络结构
class VGG11(nn.Module):
    def __init__(self):
        super(VGG11, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv5 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.conv6 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv7 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv8 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(512 * 7 * 7, 4096)
        self.dropout1 = nn.Dropout()
        self.fc2 = nn.Linear(4096, 4096)
        self.dropout2 = nn.Dropout()
        self.fc3 = nn.Linear(4096, 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 = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = self.pool3(x)
        x = F.relu(self.conv5(x))
        x = F.relu(self.conv6(x))
        x = self.pool4(x)
        x = F.relu(self.conv7(x))
        x = F.relu(self.conv8(x))
        x = self.pool5(x)
        x = x.view(-1, 512 * 7 * 7)
        x = F.relu(self.fc1(x))
        x = self.dropout1(x)
        x = F.relu(self.fc2(x))
        x = self.dropout2(x)
        x = self.fc3(x)
        return x

VGG11模型共有8个卷积层和3个全连接层。我们可以通过前向传播函数来看看输入数据从头到尾是如何依次处理的。最终输出的结果是每个类别的得分，我们可以使用softmax函数来获得每个类别的概率。

4. 训练模型

我们将使用学习率为0.001的Adam优化器和CrossEntropyLoss作为损失函数来训练VGG11模型。在每个epoch结束时，我们将会在验证集上进行测试，并记录训练和验证集上的损失和准确率。我们将训练模型200个epoch。

import torch.optim as optim

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#初始化模型、损失函数和优化器
net = VGG11().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
#训练模型
for epoch in range(200):  # 多次循环训练数据集
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 获取输入并加到GPU上
        inputs, labels = data[0].to(device), data[1].to(device)
        # 梯度清零
        optimizer.zero_grad()
        # 前向传递、后向传递及优化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        # 输出统计
        running_loss += loss.item()
        if i % 2000 == 1999:    # 每个2000批次打印一次状态
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

我们可以看到每个20个batch会输出一次loss值。由于训练的轮数比较多，如果训练时每一个batch都输出一次，会导致输出非常的频繁，因此我们将每20个batch作为一个步长。

5. 保存和加载模型

训练模型需要花费大量的时间和计算资源，且一旦训练好的模型丢失，需要重新训练模型。因此，我们需要将训练好的模型保存下来，以便之后使用。pytorch提供了一个简单的API来保存和加载训练好的模型。

#保存模型

PATH = './cifar_vgg11.pth'
torch.save(net.state_dict(), PATH)
#加载模型
net = VGG11()
net.load_state_dict(torch.load(PATH))
net.eval()

6. 模型评估

在模型评估过程中，我们将测试集的输出与标签进行比较，并计算模型的准确率。

correct = 0

total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))

我们可以看到，经过200个epoch的训练，VGG11模型在cifar10数据集上能够获得约71%的准确率。

7. 单张图像的预测

在单张图像的预测中，我们需要手动输入要进行预测的图像，并将该图像输入到训练好的VGG11模型中进行预测。

import matplotlib.pyplot as plt

import numpy as np
from PIL import Image
#将RGB值映射到[0,1]范围内
def normalize(image):
    image = np.array(image)
    return (image - np.min(image)) / (np.max(image) - np.min(image))
#测试单张图片
classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')
#加载一张图像
image_path = './cat.jpg'
image = Image.open(image_path)
image = transform(image)
image = image.unsqueeze(0).to(device)
#显示原始图像
plt.imshow(normalize(image.squeeze(0).cpu().numpy().transpose(1, 2, 0)))
plt.title('Original Image')
plt.axis('off')
plt.show()
#预测该图片所属的类别
net.eval()
outputs = net(image)
outputs = outputs / 0.6 #temperature=0.6
prob = F.softmax(outputs, dim=1)
_, predicted = torch.max(outputs.data, 1)
print('Prediction:', classes[predicted.item()])

我们可以看到，当我们输入一张猫的图片时，VGG11模型将其预测为cat。

本文介绍了使用pytorch构建VGG11模型来进行cifar10数据集的分类。我们使用了pytorch的数据预处理模块来加载cifar10数据集，使用torchvision.models模块提供的VGG11模型来进行训练和预测。我们通过训练和预测单张输入图，来展示VGG11模型在cifar10数据集上的性能。此外，我们还介绍了模型保存和加载的方法，以及单张图像预测的方法。