PyTorch: 基于【VGG16】处理MNIST数据集的图像分类任务【准确率98.9%+】

引言

在本博客中，小编将向大家介绍如何使用VGG16处理MNIST数据集的图像分类任务。MNIST数据集是一个常用的手写数字分类数据集，包含60,000个训练样本和10,000个测试样本。我们将使用Python编程语言和PyTorch深度学习框架来实现这个任务。

在Conda虚拟环境下安装pytorch

# CUDA 11.6
pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116
# CUDA 11.3
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113
# CUDA 10.2
pip install torch==1.12.1+cu102 torchvision==0.13.1+cu102 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu102
# CPU only
pip install torch==1.12.1+cpu torchvision==0.13.1+cpu torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cpu

步骤一：利用代码自动下载mnist数据集

import torchvision.datasets as datasets  
import torchvision.transforms as transforms  # 定义数据预处理操作  
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)

步骤二：搭建基于VGG16的图像分类模型

class VGGClassifier(nn.Module):def __init__(self, num_classes):super(VGGClassifier, self).__init__()self.features = models.vgg16(pretrained=True).features  # 使用预训练的VGG16模型作为特征提取器# 重构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, 4096),  # 添加一个全连接层，输入特征维度为512x7x7，输出维度为4096nn.ReLU(True),nn.Dropout(), # 随机将一些神经元“关闭”，这样可以有效地防止过拟合。nn.Linear(4096, 4096),  # 添加一个全连接层，输入和输出维度都为4096nn.ReLU(True),nn.Dropout(),nn.Linear(4096, 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 xdef _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)

步骤三：训练模型

import torch.optim as optim  
from torch.utils.data import DataLoader  # 定义超参数和训练参数  
batch_size = 64  # 批处理大小  
num_epochs = 5  # 训练轮数
learning_rate = 0.01  # 学习率
num_classes = 10  # 类别数（MNIST数据集有10个类别）  
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 判断是否使用GPU进行训练，如果有GPU则使用GPU进行训练，否则使用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))  # 打印出模型的准确率。

运行结果

在这里插入图片描述

后续模型的优化和改进建议

数据增强：通过旋转、缩放、平移等方式来增加训练数据，从而让模型拥有更好的泛化能力。
调整模型参数：可以尝试调整模型的参数，比如学习率、批次大小、迭代次数等，来提高模型的性能。
更换网络结构：可以尝试使用更深的网络结构，如ResNet、DenseNet等，来提高模型的性能。
调整优化器：本次代码采用SGD优化器，但仍可以尝试使用不同的优化器，如Adam、RMSprop等，来找到最适合我们模型的优化器。
添加正则化操作：为了防止过拟合，可以添加一些正则化项，如L1正则化、L2正则化等。
代码目前只有等训练完全结束后才能进入测试阶段，后续可以在每个epoch结束，甚至是指定的迭代次数完成后便进入测试阶段。因为训练完全结束的模型很可能已经过拟合，在测试集上不能表现较强的泛化能力。

完整代码

import torch
import torch.nn as nnimport torch.optim as optim
from torch.utils.data import DataLoaderimport torchvision.models as models
import torchvision.datasets as datasets
import torchvision.transforms as transformsimport 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, 4096),  # 添加一个全连接层，输入特征维度为512x7x7，输出维度为4096nn.ReLU(True),nn.Dropout(), #　随机将一些神经元“关闭”，有效地防止过拟合。nn.Linear(4096, 4096),  # 添加一个全连接层，输入和输出维度都为4096nn.ReLU(True),nn.Dropout(),nn.Linear(4096, 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 xdef _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 = 64  # 批处理大小
num_epochs = 5  # 训练轮数（epoch）
learning_rate = 0.01  # 学习率（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))  # 打印出模型的准确率。