基于Pytorch和Vgg16实现图片分类

2024-03-08 09:30

本文主要是介绍基于Pytorch和Vgg16实现图片分类,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

最近在基于Pytorch框架补一些CNN的基础知识,学会自己写简单的卷积神经网络,从加载数据集到训练模型、测试模型、保存模型和输出测试结果,现在来总结一下。

首先基于Pytorch实现Vgg16网络,命名为model.py(可为其他任意名字,但是后续导入时要记得更改)

import torch
import torch.nn as nnclass VGG16(nn.Module):def __init__(self):super(VGG16, self).__init__()self.layer1 = nn.Sequential(nn.Conv2d(3, 64, 3, 1, 1),nn.BatchNorm2d(64),nn.ReLU(),nn.Conv2d(64, 64, 3, 1, 1),nn.BatchNorm2d(64),nn.ReLU(),nn.MaxPool2d(2, 2))self.layer2 = nn.Sequential(nn.Conv2d(64, 128, 3, 1, 1),nn.BatchNorm2d(128),nn.ReLU(),nn.Conv2d(128, 128, 3, 1, 1),nn.BatchNorm2d(128),nn.ReLU(),nn.MaxPool2d(2, 2))self.layer3 = nn.Sequential(nn.Conv2d(128, 256, 3, 1, 1),nn.BatchNorm2d(256),nn.ReLU(),nn.Conv2d(256, 256, 3, 1, 1),nn.BatchNorm2d(256),nn.ReLU(),nn.Conv2d(256, 256, 3, 1, 1),nn.BatchNorm2d(256),nn.ReLU(),nn.MaxPool2d(2, 2))self.layer4 = nn.Sequential(nn.Conv2d(256, 512, 3, 1, 1),nn.BatchNorm2d(512),nn.ReLU(),nn.Conv2d(512, 512, 3, 1, 1),nn.BatchNorm2d(512),nn.ReLU(),nn.Conv2d(512, 512, 3, 1, 1),nn.BatchNorm2d(512),nn.ReLU(),nn.MaxPool2d(2, 2))self.layer5 = nn.Sequential(nn.Conv2d(512, 512, 3, 1, 1),nn.BatchNorm2d(512),nn.ReLU(),nn.Conv2d(512, 512, 3, 1, 1),nn.BatchNorm2d(512),nn.ReLU(),nn.Conv2d(512, 512, 3, 1, 1),nn.BatchNorm2d(512),nn.ReLU(),nn.MaxPool2d(2, 2))self.fc1 = nn.Sequential(nn.Flatten(),nn.Linear(512, 512),nn.ReLU(),nn.Dropout())self.fc2 = nn.Sequential(nn.Linear(512, 256),nn.ReLU(),nn.Dropout())self.fc3 = nn.Linear(256, 10)def forward(self, x):x = self.layer1(x)x = self.layer2(x)x = self.layer3(x)x = self.layer4(x)x = self.layer5(x)x = self.fc1(x)x = self.fc2(x)x = self.fc3(x)return xif __name__ == '__main__':VGG16 = VGG16()input = torch.ones((64, 3, 32, 32))output = VGG16(input)print(output.shape)

然后编写训练文件,导入需要的库(有些库没有的话需要提前安装)。

import torchvision
from torch import optim
from torch.utils.data import DataLoader
import torch.nn as nn
from model import *
import matplotlib.pyplot as plt
import time

下载CIFAR10数据集,并设置batch_size,第一次运行时会自动下载,之后直接加载已经下载好的。当然也可以使用其他数据集,从官方文档https://pytorch.org/vision/stable/datasets.html可以看到有各种用于分类、分割、检测等其他视觉任务的数据集,但是为了方便学习,这里选择较小的CIFAR10数据集。

rain_data = torchvision.datasets.CIFAR10(root="data", train=True, transform=torchvision.transforms.ToTensor(),download=True)
test_data = torchvision.datasets.CIFAR10(root="data", train=False, transform=torchvision.transforms.ToTensor(),download=True)
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

创建记录文档,名字为当下时间,将训练过程和结果保存并打印出来。

t = time.strftime('%Y-%m-%d_%H-%M-%S')  # 切记中间符号不能用冒号,不然会报错
file = open('logs/{}.txt'.format(t), 'w')
train_data_size = len(train_data)
test_data_size = len(test_data)
print('The size of train_data:{}'.format(train_data_size))
file.write('The size of train_data:{}'.format(train_data_size)+'\n')
print('The size of train_data:{}'.format(test_data_size))
file.write('The size of train_data:{}'.format(test_data_size)+'\n')

创建一些中间参数,配置损失函数和优化器,也可从官方文档中选择其他损失函数torch.nn.functional — PyTorch 1.11.0 documentation和优化器torch.optim — PyTorch 1.11.0 documentation

epoch = 10
train_step = 0
test_step = 0vgg16 = VGG16()  # 实例化网络
loss_fn = nn.CrossEntropyLoss()  learning_rate = 1e-2
# optimizer = optim.SGD(vgg16.parameters(), lr=learning_rate, momentum=0.9)
optimizer = optim.Adam(vgg16.parameters(), lr=learning_rate)losses = []  # 用于存储损失值,便于后面画损失变化曲线图

开始训练

for i in range(epoch):print('-----epoch{}-----'.format(i))file.write('-----epoch{}-----'.format(i)+'\n')# 训练步骤开始vgg16.train()for data in train_dataloader:images, targets = dataoutput = vgg16(images)loss = loss_fn(output, targets)# 优化器优化模型optimizer.zero_grad()loss.backward()optimizer.step()train_step += 1if train_step % 50 == 0:print('train_step:{},loss:{}'.format(train_step, loss.item()))file.write('train_step:{},loss:{}'.format(train_step, loss.item())+'\n')losses.append(loss.item())# 测试步骤开始vgg16.eval()accuracy = 0with torch.no_grad():for data in test_dataloader:images, targets = dataoutput = vgg16(images)current_acc = (output.argmax(1) == targets).sum()accuracy += current_accacc = accuracy / test_data_sizeprint('-------eval accuracy:{}'.format(acc))file.write('-------eval accuracy:{}'.format(acc)+'\n')torch.save(tudui, 'checkpoints/vgg16_{}.pth'.format(i))print('The model has been saved')file.write('The model has been saved')file.close()

保存的记录文档和打印输出如下所示

 

画出损失曲线变化图

plt.subplot(1, 1, 1)
plt.plot(losses, label='Training Loss')
plt.title('Training Loss')
plt.legend()
plt.show()

 编写测试文件,将训练好的模型用于预测任意图片的类别。

import torch
from PIL import Image
from torchvision import transformsimage_path = 'imgs/dog.jpg'  # 随意选择一张图片进行测试
image = Image.open(image_path)
image = image.convert('RGB')
transform = transforms.Compose([transforms.Resize(32, 32), transforms.ToTensor()])  # 将图片大小改为32×32,并转换为Tensor类型
image = transform(image)classes = ['plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']model = torch.load('checkpoints/vgg16_30.pth')  # 加载训练好的模型image = torch.reshape(image, (1, 3, 32, 32))
model.eval()
with torch.no_grad():output = model(image)
index = output.argmax(1)
pred = classes[index]
print(pred)

                                 

这篇关于基于Pytorch和Vgg16实现图片分类的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/786684

相关文章

关于集合与数组转换实现方法

《关于集合与数组转换实现方法》:本文主要介绍关于集合与数组转换实现方法,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教... 目录1、Arrays.asList()1.1、方法作用1.2、内部实现1.3、修改元素的影响1.4、注意事项2、list.toArray()2.1、方

使用Python实现可恢复式多线程下载器

《使用Python实现可恢复式多线程下载器》在数字时代,大文件下载已成为日常操作,本文将手把手教你用Python打造专业级下载器,实现断点续传,多线程加速,速度限制等功能,感兴趣的小伙伴可以了解下... 目录一、智能续传:从崩溃边缘抢救进度二、多线程加速:榨干网络带宽三、速度控制:做网络的好邻居四、终端交互

java实现docker镜像上传到harbor仓库的方式

《java实现docker镜像上传到harbor仓库的方式》:本文主要介绍java实现docker镜像上传到harbor仓库的方式,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地... 目录1. 前 言2. 编写工具类2.1 引入依赖包2.2 使用当前服务器的docker环境推送镜像2.2

C++20管道运算符的实现示例

《C++20管道运算符的实现示例》本文简要介绍C++20管道运算符的使用与实现,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧... 目录标准库的管道运算符使用自己实现类似的管道运算符我们不打算介绍太多,因为它实际属于c++20最为重要的

Java easyExcel实现导入多sheet的Excel

《JavaeasyExcel实现导入多sheet的Excel》这篇文章主要为大家详细介绍了如何使用JavaeasyExcel实现导入多sheet的Excel,文中的示例代码讲解详细,感兴趣的小伙伴可... 目录1.官网2.Excel样式3.代码1.官网easyExcel官网2.Excel样式3.代码

python实现对数据公钥加密与私钥解密

《python实现对数据公钥加密与私钥解密》这篇文章主要为大家详细介绍了如何使用python实现对数据公钥加密与私钥解密,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下... 目录公钥私钥的生成使用公钥加密使用私钥解密公钥私钥的生成这一部分,使用python生成公钥与私钥,然后保存在两个文

浏览器插件cursor实现自动注册、续杯的详细过程

《浏览器插件cursor实现自动注册、续杯的详细过程》Cursor简易注册助手脚本通过自动化邮箱填写和验证码获取流程,大大简化了Cursor的注册过程,它不仅提高了注册效率,还通过友好的用户界面和详细... 目录前言功能概述使用方法安装脚本使用流程邮箱输入页面验证码页面实战演示技术实现核心功能实现1. 随机

Golang如何对cron进行二次封装实现指定时间执行定时任务

《Golang如何对cron进行二次封装实现指定时间执行定时任务》:本文主要介绍Golang如何对cron进行二次封装实现指定时间执行定时任务问题,具有很好的参考价值,希望对大家有所帮助,如有错误... 目录背景cron库下载代码示例【1】结构体定义【2】定时任务开启【3】使用示例【4】控制台输出总结背景

Golang如何用gorm实现分页的功能

《Golang如何用gorm实现分页的功能》:本文主要介绍Golang如何用gorm实现分页的功能方式,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教... 目录背景go库下载初始化数据【1】建表【2】插入数据【3】查看数据4、代码示例【1】gorm结构体定义【2】分页结构体

在Golang中实现定时任务的几种高效方法

《在Golang中实现定时任务的几种高效方法》本文将详细介绍在Golang中实现定时任务的几种高效方法,包括time包中的Ticker和Timer、第三方库cron的使用,以及基于channel和go... 目录背景介绍目的和范围预期读者文档结构概述术语表核心概念与联系故事引入核心概念解释核心概念之间的关系