VGG16模型实现MNIST图像分类

MNIST图像数据集

MNIST（Modified National Institute of Standards and Technology）是一个经典的机器学习数据集，常用于训练和测试图像处理和机器学习算法，特别是在数字识别领域。该数据集包含了大约 7 万张手写数字图片，其中 6 万张是用于训练，1 万张用于测试。每张图片都是 28x28 像素的灰度图像，展示了从 0 到 9 的手写数字。这些图像已经被处理过，以使得数字在图像中居中且尺寸一致。

MNIST 数据集是一个广泛被用于测试新的机器学习算法的基准，因为它相对较小，易于理解，且可以用于快速验证算法的有效性。许多人使用 MNIST 作为开始学习深度学习的入门数据集，因为它提供了一个简单但具有挑战性的任务，即将手写数字图像分类为相应的数字。

尽管 MNIST 已经存在了很长时间，但它仍然是一个重要的基准数据集，特别是对于新的机器学习研究和算法的初步测试。MINIST数据集中部分图片如下所示：

下载MNIST数据集

由于MINIST作为经典数据集，已经被内嵌在torchvision库中的dataset中了，所以直接使用代码datasets.MNIST进行下载即可。

下载后的文件格式如下图所示。

搭建VGG16图像分类模型

class VGGClassifier(nn.Module):
    def __init__(self, num_classes):
        super(VGGClassifier, self).__init__()
        self.features = models.vgg16(pretrained=True).features  # 使用预训练的VGG16模型作为特征提取器
        # 重构网络的第一层卷积层，适配mnist数据的灰度图像格式
        self.features[0] = nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 256),  # 添加一个全连接层，输入特征维度为512x7x7，输出维度为4096
            nn.ReLU(True),
            nn.Dropout(), #　随机将一些神经元“关闭”，有效地防止过拟合。
            nn.Linear(256, 256),  # 添加一个全连接层，输入和输出维度都为4096
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(256, num_classes),  # 添加一个全连接层，输入维度为4096，输出维度为类别数（10）
        )
        self._initialize_weights()  # 初始化权重参数

定义VGG网络结构如上所示，在上面代码中我定义了一个基于 VGG16 架构分类器的模型。VGG16 是一种经典的卷积神经网络模型，由 16 层深度的卷积层和全连接层组成，所构建的 VGGClassifier 类的网络结构包含两个主要部分：

特征提取器（features）：这部分使用了预训练的 VGG16 模型的特征提取器。通过调用 models.vgg16(pretrained=True).features 来加载 VGG16 的特征提取器部分。然后，将第一层卷积层的输入通道数从 3 修改为 1，以适应 MNIST 数据集的灰度图像格式。

分类器（classifier）：这部分是自定义的分类器，用于对提取的特征进行分类。首先，通过几个全连接层将特征图展平成一维张量，然后通过一系列的线性层和激活函数对特征进行处理。具体来说，包括：一个包含 256 个神经元的全连接层，输入维度为 512x7x7（经过 VGG16 的特征提取器后的输出尺寸），使用 ReLU 激活函数。一个 Dropout 层，用于防止过拟合，随机关闭一些神经元。一个包含 256 个神经元的全连接层，使用 ReLU 激活函数。再次添加一个 Dropout 层。最后是一个包含 num_classes 个神经元的全连接层，用于输出最终的类别预测结果。

通过上述方式，整个网络结构将 VGG16 的特征提取器和自定义的分类器相结合，以适应 MNIST 数据集的图像分类任务。

构建的VGG网络结构如下图所示：

VGG网络结构图

模型训练

# 定义超参数和训练参数
batch_size = 16  # 批处理大小
num_epochs = 5  # 训练轮数（epoch）
learning_rate = 0.001  # 学习率（learning rate）
num_classes = 10  # 类别数（MNIST数据集有10个类别）
device = torch.device(
    "cuda:0" if torch.cuda.is_available() else "cpu")  # 判断是否使用GPU进行训练，如果有GPU则使用第一个GPU（cuda:0）进行训练，否则使用CPU进行训练。

模型参数设置如下表所示（代码见上）

模型超参数	数值
batchsize	16
num_epochs	5
learning_rate	0.001
num_classes	10

由于MINIST数据集样本数量较大，所以对于上述代码训练速度也会较慢，我考虑使用我的笔记本电脑独显进行运算，却发现电脑显存不够，于是我调小batchsize与epoch，并降低学习率learning rate才让GPU勉强能够运行上面代码，并获得到了模型model.pth，最终获得模型在测试集上面的识别精度为96.7%，精度还是比较高的。（由于笔记本电脑性能有限，在处理较大规模数据的小型项目时速度较慢，故上述代码运行了一下午左右的时间才跑完）。

模型测试

使用上面模型进行手写数字识别的检验。绘制一张图片上面含有9张子图，随机选取识别结果的9张进行展示 。识别效果以及运行结果如下图所示。

 

附录：

 VGG训练代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.models as models
import torchvision.datasets as datasets
import torchvision.transforms as transforms

import warnings
warnings.filterwarnings("ignore")

# 定义数据预处理操作
transform = transforms.Compose([
    transforms.Resize(224), # 将图像大小调整为(224, 224)
    transforms.ToTensor(),  # 将图像转换为PyTorch张量
    transforms.Normalize((0.5,), (0.5,))  # 对图像进行归一化
])

# 下载并加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)


class VGGClassifier(nn.Module):
    def __init__(self, num_classes):
        super(VGGClassifier, self).__init__()
        self.features = models.vgg16(pretrained=True).features  # 使用预训练的VGG16模型作为特征提取器
        # 重构网络的第一层卷积层，适配mnist数据的灰度图像格式
        self.features[0] = nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 256),  # 添加一个全连接层，输入特征维度为512x7x7，输出维度为4096
            nn.ReLU(True),
            nn.Dropout(), #　随机将一些神经元“关闭”，有效地防止过拟合。
            nn.Linear(256, 256),  # 添加一个全连接层，输入和输出维度都为4096
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(256, num_classes),  # 添加一个全连接层，输入维度为4096，输出维度为类别数（10）
        )
        self._initialize_weights()  # 初始化权重参数

    def forward(self, x):
        x = self.features(x)  # 通过特征提取器提取特征
        x = x.view(x.size(0), -1)  # 将特征张量展平为一维向量
        x = self.classifier(x)  # 通过分类器进行分类预测
        return x

    def _initialize_weights(self):  # 定义初始化权重的方法，使用Xavier初始化方法
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

# 定义超参数和训练参数
batch_size = 16  # 批处理大小
num_epochs = 5  # 训练轮数（epoch）
learning_rate = 0.001  # 学习率（learning rate）
num_classes = 10  # 类别数（MNIST数据集有10个类别）
device = torch.device(
    "cuda:0" if torch.cuda.is_available() else "cpu")  # 判断是否使用GPU进行训练，如果有GPU则使用第一个GPU（cuda:0）进行训练，否则使用CPU进行训练。

# 定义数据加载器
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

# 初始化模型和优化器
model = VGGClassifier(num_classes=num_classes).to(device)  # 将模型移动到指定设备（GPU或CPU）
criterion = nn.CrossEntropyLoss()  # 使用交叉熵损失函数
optimizer = optim.SGD(model.parameters(), lr=learning_rate)  # 使用随机梯度下降优化器（SGD）

# 训练模型
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)  # 将图像数据移动到指定设备
        labels = labels.to(device)  # 将标签数据移动到指定设备

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()  # 清空梯度缓存
        loss.backward()  # 计算梯度
        optimizer.step()  # 更新权重参数

        if (i + 1) % 100 == 0:  # 每100个batch打印一次训练信息
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, num_epochs, i + 1, len(train_loader),
                                                                     loss.item()))

# 训练结束，保存模型参数
torch.save(model.state_dict(), './model.pth')

# 加载训练好的模型参数
model.load_state_dict(torch.load('./model.pth'))
model.eval()  # 将模型设置为评估模式，关闭dropout等操作

# 定义评估指标变量
correct = 0  # 记录预测正确的样本数量
total = 0  # 记录总样本数量

# 测试模型性能
with torch.no_grad():  # 关闭梯度计算，节省内存空间
    for images, labels in test_loader:
        images = images.to(device)  # 将图像数据移动到指定设备
        labels = labels.to(device)  # 将标签数据移动到指定设备
        outputs = model(images)  # 模型前向传播，得到预测结果
        _, predicted = torch.max(outputs.data, 1)  # 取预测结果的最大值对应的类别作为预测类别
        total += labels.size(0)  # 更新总样本数量
        correct += (predicted == labels).sum().item()  # 统计预测正确的样本数量

# 计算模型准确率并打印出来
accuracy = 100 * correct / total  # 计算准确率，将正确预测的样本数量除以总样本数量并乘以100得到百分比形式的准确率。
print('Accuracy of the model on the test images: {} %'.format(accuracy))  # 打印出模型的准确率。