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))  # 打印出模型的准确率。