博客

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

  • Post category:Python

PyTorch VGG11识别CIFAR10数据集

本文将详细讲解如何使用PyTorch的VGG11模型对CIFAR10数据集进行分类，并提供训练和预测单张输入图片的操作。

准备工作

在开始之前，需要安装PyTorch和CIFAR10数据。可以使用以下命令来安装：

pip install torch torchvision

CIFAR10数据集可以在PyTorch中直接下载，也可以从官网下载并手动导入。在本文中，我们将使用PyTorch中的自动下载功能。

加载数据集

首先，我们需要加载CIFAR10数据集。可以使用以下代码来加载数据集：

import torch
import torchvision
import torchvision.transforms as transforms

# 定义数据预处理
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=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)

# 加载测试集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

# 定义类别标签
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

在上面的代码中，我们使用transforms.Compose定义了数据预处理，包括将图像转换为张量并进行归一化。然后，我们使用torchvision.datasets.CIFAR10加载训练集和测试集，并使用torch.utils.data.DataLoader定义了数据加载器。最后，我们定义了类别标签。

定义模型

接下来，我们需要定义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.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.conv5 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.conv6 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv7 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv8 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.fc = nn.Linear(512 * 2 * 2, 4096)
        self.fc2 = nn.Linear(4096, 4096)
        self.fc3 = nn.Linear(4096, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv5(x))
        x = F.relu(self.conv6(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv7(x))
        x = F.relu(self.conv8(x))
        x = F.max_pool2d(x, 2)
        x = x.view(-1, 512 * 2 * 2)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = F.relu(self.fc2(x))
        x = F.dropout(x, training=self.training)
        x = self.fc3(x)
        return x

# 实例化模型
net = VGG11()

在上面的代码中，我们定义了VGG11模型，并实例化了模型。

训练模型

接下来，我们需要训练模型。可以使用以下代码来训练模型：

import torch.optim as optim

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 训练模型
for epoch in range(2):  #次循环数据集

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 获取输入
        inputs, labels = data

        # 梯度清零
        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

print('Finished Training')

在上面的代码中，我们定义了损失函数和优化器，并使用for循环多次循环数据集进行训练。在每个小批量数据上，我们执行正向传播、反向传播和优化，并打印统计信息。

预测单张输入图片

最后，我们需要使用训练好的模型对单张输入图片进行预测。可以使用以下代码来预测单张输入图片：

import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

# 加载单张输入图片
img = Image.open('test.jpg')
img = transform(img)
img = img.unsqueeze(0)

# 预测单张输入
outputs = net(img)
_, predicted = torch.max(outputs, 1)

# 打印预测结果
print('Predicted: ', classes[predicted[0]])

# 显示输入图片
img = img.squeeze()
img = img.numpy().transpose((1, 2, 0))
img = img * 0.5 + 0.5
plt.imshow(img)
plt.show()

在上面的代码中，我们使用PIL库加载单张输入图片，并使用transform进行数据预处理。然后，我们使用训练好的模型对输入图片进行预测，并打印预测结果。最后，我们使用matplotlib库显示输入图片。

完整示例

下面是一个完整的示例，包括加载数据集、定义模型、训练模型和预测单张输入图片：

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

# 定义数据预处理
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=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)

# 加载测试集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

# 定义类别标签
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# 定义VGG11模型
class VGG11(nn.Module):
    def __init__(self):
        super(VGG11, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.conv5 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.conv6 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv7 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv8 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(512 * 2 * 2, 4096)
        self.fc2 = nn.Linear(4096, 4096)
        self.fc3 = nn.Linear(4096, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv5(x))
        x = F.relu(self.conv6(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv7(x))
        x = F.relu(self.conv8(x))
        x = F.max_pool2d(x, 2)
        x = x.view(-1, 512 * 2 * 2)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = F.relu(self.fc2(x))
        x = F.dropout(x, training=self.training)
        x = self.fc3(x)
        return x

# 实例化模型
net = VGG11()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 训练模型
for epoch in range(2):  #次循环数据集

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 获取输入
        inputs, labels = data

        # 梯度清零
        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

print('Finished Training')

# 加载单张输入图片
img = Image.open('test.jpg')
img = transform(img)
img = img.unsqueeze(0)

# 预测单张输入
outputs = net(img)
_, predicted = torch.max(outputs, 1)

# 打印预测结果
print('Predicted: ', classes[predicted[0]])

# 显示输入图片
img = img.squeeze()
img = img.numpy().transpose((1, 2, 0))
img = img * 0.5 + 0.5
plt.imshow(img)
plt.show()

在上面的代码中，我们首先加载数据集，然后定义了VGG11模型，并实例化了模型。接着，我们定义了损失函数和优化器，并使用for循环多次循环数据集进行训练。在每个小批量数据上，我们执行正向传播、反向传播和优化，并打印统计信息。最后，我们使用PIL库加载单张输入图片，并使用训练好的模型对输入图片进行预测，并打印预测结果。最后，我们使用matplotlib库显示输入图片。

总结

通过本文的学习，您可以了解如何使用PyTorch构建和训练深度学习模型，并对图像分类任务有更深入的理解。

Tags: