引言
在本博客中，小编将向大家介绍如何使用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，输出维度为4096
nn.ReLU(True),
nn.Dropout(), # 随机将一些神经元“关闭”，这样可以有效地防止过拟合。
nn.Linear(4096, 4096), # 添加一个全连接层，输入和输出维度都为4096
nn.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 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)
步骤三：训练模型
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 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, 4096), # 添加一个全连接层，输入特征维度为512x7x7，输出维度为4096
nn.ReLU(True),
nn.Dropout(), # 随机将一些神经元“关闭”，有效地防止过拟合。
nn.Linear(4096, 4096), # 添加一个全连接层，输入和输出维度都为4096
nn.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 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 = 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)) # 打印出模型的准确率。
结束语
如果本博文对你有所帮助/启发，可以点个赞/收藏支持一下，如果能够持续关注，小编感激不尽~
如果有相关需求/问题需要小编帮助，欢迎私信~
小编会坚持创作，持续优化博文质量，给读者带来更好de阅读体验~
