PyTorch: 基于VGG16处理MNIST数据集的图像分类任务

目录

  • 引言
  • 在Conda虚拟环境下安装pytorch
  • 步骤一：利用代码自动下载mnist数据集
  • 步骤二：搭建基于VGG16的图像分类模型
  • 步骤三：训练模型
  • 步骤四：测试模型
  • 运行结果
  • 后续模型的优化和改进建议
  • 完整代码
  • 结束语

引言

在本博客中，小编将向大家介绍如何使用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))  # 打印出模型的准确率。

运行结果

在这里插入图片描述

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

  1. 数据增强：通过旋转、缩放、平移等方式来增加训练数据，从而让模型拥有更好的泛化能力。
  2. 调整模型参数：可以尝试调整模型的参数，比如学习率、批次大小、迭代次数等，来提高模型的性能。
  3. 更换网络结构：可以尝试使用更深的网络结构，如ResNet、DenseNet等，来提高模型的性能。
  4. 调整优化器：本次代码采用SGD优化器，但仍可以尝试使用不同的优化器，如Adam、RMSprop等，来找到最适合我们模型的优化器。
  5. 添加正则化操作：为了防止过拟合，可以添加一些正则化项，如L1正则化、L2正则化等。
  6. 代码目前只有等训练完全结束后才能进入测试阶段，后续可以在每个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))  # 打印出模型的准确率。

结束语

如果本博文对你有所帮助/启发，可以点个赞/收藏支持一下，如果能够持续关注，小编感激不尽~
如果有相关需求/问题需要小编帮助，欢迎私信~
小编会坚持创作，持续优化博文质量，给读者带来更好de阅读体验~

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：http://www.hqwc.cn/news/268394.html

如若内容造成侵权/违法违规/事实不符，请联系编程知识网进行投诉反馈email:809451989@qq.com，一经查实，立即删除！

相关文章

猎豹浏览器如何设置ip使用？socks5在网络安全中有什么优势？

猎豹浏览器如何设置ip使用？socks5在网络安全中有什么优势？

猎豹浏览器如何设置ip使用？socks在网络安全中有什么优势？ 一、猎豹浏览器如何设置ip使用？ 在使用猎豹浏览器时，可以通过以下步骤来设置IP使用： 1. 打开猎豹浏览器，点击右上角的“菜单”按钮，在…
阅读更多...
软文是如何连接品牌与用户的？媒介盒子为您解答

软文是如何连接品牌与用户的？媒介盒子为您解答

软文作为品牌向用户传递信息或价值传递的重要媒介，在品牌营销中的场景中发挥着重要作用。那么商家如何做，才能使软文成功连接品牌与用户呢？下面就让媒介盒子分享。 一、软文是什么 软文是一种带有目的的表达，软文目的大概分为四种…
阅读更多...
c++ map

c++ map

unordered_map #include <iostream> #include <string> #include <unordered_map>int main() {// 创建包含三个字符串的&#xff08;映射到字符串的&#xff09;unordered_mapstd::unordered_map<std::string, std::string> u {{"red", &qu…
阅读更多...
C语言普里姆（Prim）算法实现计算国家建设高铁运输网最低造价的建设方案

C语言普里姆（Prim）算法实现计算国家建设高铁运输网最低造价的建设方案

背景&#xff1a; 描述&#xff1a;为促进全球更好互联互通&#xff0c;亚投行拟在一带一路沿线国家建设高铁运输网&#xff0c;请查阅相关资料 画出沿线国家首都或某些代表性城市的连通图&#xff0c;为其设计长度最短或造价最低的高铁建设方案。 要求&#xff1a;抽象出的图…
阅读更多...
智能优化算法应用：基于鸡群算法3D无线传感器网络(WSN)覆盖优化 - 附代码

智能优化算法应用：基于鸡群算法3D无线传感器网络(WSN)覆盖优化 - 附代码

智能优化算法应用&#xff1a;基于鸡群算法3D无线传感器网络(WSN)覆盖优化 - 附代码 文章目录 智能优化算法应用&#xff1a;基于鸡群算法3D无线传感器网络(WSN)覆盖优化 - 附代码1.无线传感网络节点模型2.覆盖数学模型及分析3.鸡群算法4.实验参数设定5.算法结果6.参考文献7.MA…
阅读更多...
kubernetes（k8s) Yaml 文件详解

kubernetes（k8s) Yaml 文件详解

YAML格式&#xff1a;用于配置和管理&#xff0c;YAML是一种简洁的非标记性语言&#xff0c;内容格式人性化&#xff0c;较易读。 1、查看API 资源版本标签 kubectl api-versions 2、编写资源配置清单 kubectl create -f nginx-test.yaml --validatefalse 2.3 查看创建的po…
阅读更多...
讲解充电桩测试项目

讲解充电桩测试项目

充电桩测试项目是为了保证充电桩的质量和安全性&#xff0c;对充电桩进行全面、系统的检测和评估。功能测试是充电桩测试的基础&#xff0c;主要检查充电桩的各项功能是否正常。包括充电接口连接测试、充电过程控制测试、充电状态显示测试等。通过功能测试&#xff0c;可以确保…
阅读更多...
python自动化测试实战 —— WebDriver API的使用

python自动化测试实战 —— WebDriver API的使用

软件测试专栏 感兴趣可看&#xff1a;软件测试专栏 自动化测试学习部分源码 python自动化测试相关知识&#xff1a; 【如何学习Python自动化测试】—— 自动化测试环境搭建 【如何学习python自动化测试】—— 浏览器驱动的安装 以及 如何更…
阅读更多...
JIT即时编译器深度解析——Java性能提升利器

JIT即时编译器深度解析——Java性能提升利器

文章目录 一、JIT概述1、为什么要用JIT即时编译器2、C1、C2与Graal编译器3、分层编译4、热点代码5、热点探测&#xff08;1&#xff09;方法调用计数器&#xff08;2&#xff09;回边计数器 二、编译优化技术1、方法内联&#xff08;1&#xff09;什么是方法内联&#xff08;2&…
阅读更多...
如何充分准备面试，迅速融入团队并在工作中取得卓越成就

如何充分准备面试，迅速融入团队并在工作中取得卓越成就

首先&#xff0c;关于如何筹备面试&#xff0c;首先需要对所申请公司与职位进行深入的调查了解&#xff0c;并依据可能提出的面试问题预先准备相应的答案&#xff0c;并提前调试面试所需的仪器设备。同时&#xff0c;也要注重自身形象的塑造。更为关键的是 1. 在计算机领域的面…
阅读更多...
读一张AFL终端截图

读一张AFL终端截图

以下回答来源于ChatGPT plus 英文版 The image you uploaded shows a screenshot of a terminal or command-line interface running a fuzz testing program. Heres an analysis based on the visible text and layout: Program Name: The program is named "american f…
阅读更多...
java基础之抽象的概念（全网最详细）

java基础之抽象的概念（全网最详细）

一.抽象的概念 如下图所示&#xff1a; 如上图所示&#xff0c;一个父类为图形&#xff0c;计算图形的面积方法&#xff0c;那么就会有疑问&#xff0c;计算哪个图形的面积呢&#xff1f;所以这个计算图形的面积方法就称为抽象方法&#xff1b; 二.抽象方法和抽象类的格式 …
阅读更多...
推荐文章
最新文章

Copyright @ 2022~2023