神经网络/深度学习
第一章 Python机器学习入门之VGG16的使用
前言
这篇文章主要是本人在进行学习vgg-16时所进行的代码复现，将其学习记录下来。主要的数据集是UI和壁纸之间的分类预测。
一、VGG16是什么？
vgg16是由Visual Geometry Group组的Simonyan和Zisserman在文献《Very Deep Convolutional Networks for Large Scale Image Recognition》
[1]Simonyan K, Zisserman A. Very deep convolutional networks for large-scale image recognition[J]. arXiv preprint arXiv:1409.1556, 2014.
中提出卷积神经网络模型，具体的解释这个网络大家可以去看看其它博主的讲解，我在这里就不多做解释了。po一张网络图上来
二、VGG16代码的复现
1.对图片重命名
相关代码可以去看我的上一篇博客，我在这里也就不占用篇幅了。例如壁纸就叫picture.xxx.png，最后预测的二分类也是以上两个类别，文件夹则放在train文件夹下，如图。
2.准备工作
我所复现的是下面这位大佬的相关代码，大佬讲的非常详细，链接如下：
[2]https://www.bilibili.com/video/BV1X3411N7aj/?spm_id_from=333.337.search-card.all.click&vd_source=b9a1a486cbe5d7fe623135210f75aca8
- 第一步：在更改完对应的名称之后，要先进行vgg16模型的下载，可以直接运行net.py文件进行模型的下载，也可以通过下面的链接提前下载到本地
https://download.pytorch.org/models/vgg16-397923af.pth
- 第二步：下载好了之后运行txt.py文件，并在第4行修改你所需要的种类
import os
from os import getcwd
classes=['phone','picture']#在这里增加种类
sets=['train']
if __name__=='__main__':
wd=getcwd()
for se in sets:
list_file=open('cls_'+ se +'.txt','w')
datasets_path=se
types_name=os.listdir(datasets_path)#os.listdir() 方法用于返回指定的文件夹包含的文件或文件夹的名字的列表
for type_name in types_name:
if type_name not in classes:
continue
cls_id=classes.index(type_name)#输出0-1
photos_path=os.path.join(datasets_path,type_name)
photos_name=os.listdir(photos_path)
for photo_name in photos_name:
_,postfix=os.path.splitext(photo_name)#该函数用于分离文件名与拓展名
if postfix not in['.jpg','.png','.jpeg']:
continue
list_file.write(str(cls_id)+';'+'%s/%s'%(wd, os.path.join(photos_path,photo_name)))
list_file.write('\n')
list_file.close()
运行过后会出现一个名叫cls_train.txt的文件，内容大致如下：
种类0是phone（UI），种类1则是picture（壁纸）
3.训练模型
import torch
import torch.nn as nn
from net import vgg16
from torch.utils.data import DataLoader#工具取黑盒子，用函数来提取数据集中的数据（小批次）
from data import *
import matplotlib.pyplot as plt # 导入matplotlib绘图库
# 在训练开始前初始化精度列表
train_accuracies = []
val_accuracies = []
'''数据集'''
annotation_path='cls_train.txt'#读取数据集生成的文件
with open(annotation_path,'r') as f:
lines=f.readlines()
np.random.seed(10101)#函数用于生成指定随机数
np.random.shuffle(lines)#数据打乱
np.random.seed(None)
num_val=int(len(lines)*0.2)#十分之一数据用来测试
num_train=len(lines)-num_val
#输入图像大小
input_shape=[224,224] #导入图像大小
train_data=DataGenerator(lines[:num_train],input_shape,True)
val_data=DataGenerator(lines[num_train:],input_shape,False)
val_len=len(val_data)
print(val_len)#返回测试集长度
# 取黑盒子工具
"""加载数据"""
gen_train=DataLoader(train_data,batch_size=4)#训练集batch_size读取小样本，规定每次取多少样本
gen_test=DataLoader(val_data,batch_size=4)#测试集读取小样本
'''构建网络'''
device=torch.device('cuda'if torch.cuda.is_available() else "cpu")#电脑主机的选择
net=vgg16(True, progress=True,num_classes=2)#定于分类的类别
net.to(device)
'''选择优化器和学习率的调整方法'''
lr=0.0001#定义学习率
optim=torch.optim.Adam(net.parameters(),lr=lr)#导入网络和学习率
sculer=torch.optim.lr_scheduler.StepLR(optim,step_size=1)#步长为1的读取
'''训练'''
epochs=2#读取数据次数，每次读取顺序方式不同
for epoch in range(epochs):
total_train=0 #定义总损失
for data in gen_train:
img,label=data
with torch.no_grad():
img =img.to(device)
label=label.to(device)
optim.zero_grad()
output=net(img)
train_loss=nn.CrossEntropyLoss()(output,label).to(device)
train_loss.backward()#反向传播
optim.step()#优化器更新
total_train+=train_loss #损失相加
sculer.step()
total_test=0#总损失
total_accuracy=0#总精度
for data in gen_test:
img,label =data #图片转数据
with torch.no_grad():
img=img.to(device)
label=label.to(device)
optim.zero_grad()#梯度清零
out=net(img)#投入网络
test_loss=nn.CrossEntropyLoss()(out,label).to(device)
total_test+=test_loss#测试损失，无反向传播
accuracy=((out.argmax(1)==label).sum()).clone().detach().cpu().numpy()#正确预测的总和比测试集的长度，即预测正确的精度
total_accuracy+=accuracy
# 假设你已经计算了当前轮次的验证精度并存储在变量`val_accuracy`中
val_accuracy_percent = (total_accuracy / val_len) * 100 # 将精度转换为百分比
val_accuracies.append(val_accuracy_percent)
# 打印当前轮次的训练和验证损失、精度等信息
print(f"Epoch {epoch + 1}/{epochs}")
print("训练集上的损失：{}".format(total_train))
print("测试集上的损失：{}".format(total_test))
print("测试集上的精度：{:.1%}".format(total_accuracy/val_len))#百分数精度，正确预测的总和比测试集的长度
torch.save(net.state_dict(), "fenlei{}.pth".format(epoch + 1)) # 模型的名称
print("模型已保存")
# 所有训练轮次完成后，绘制精度变化图
plt.figure(figsize=(10, 6))
plt.plot(range(1, epochs + 1), val_accuracies, label='Validation Accuracy')
plt.title('Accuracy over Epochs')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.grid(True)
# 在显示图像之前保存它
plt.savefig('accuracy_plot.png') # 将图像保存到当前目录下的accuracy_plot.png文件中
plt.show() # 然后显示图像
这是mian.py也就是训练代码，其中规定了各种参数，包括训练轮次，学习率等。
4.进行预测
预测的代码自己进行了一定的修改，主要是讲显示格式改成了json，其他的基本没动。
from torchvision import transforms
from PIL import Image
import torch
import torch.nn.functional as F
from net import vgg16 # 确保您的目录中有这个文件和相应的VGG16模型实现
import json
# 设置可以检测的图像路径
test_pth = '图片路径'
test = Image.open(test_pth)
# 处理图片
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.Grayscale(num_output_channels=3), # 将图像转换为三通道
transforms.ToTensor()
])
image = transform(test)
# 加载网络
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 选择CPU或GPU
net = vgg16() # 初始化网络
model = torch.load(r"训练好的模型位置", map_location=device) # 加载已训练的模型权重
net.load_state_dict(model) # 将权重加载到网络
net.eval() # 设置为评估模式
# 调整图片维度以符合网络输入
image = torch.reshape(image, (1, 3, 224, 224))
# 进行预测
with torch.no_grad():
out = net(image)
# 使用softmax获取概率分布，并找到最大概率值
out = F.softmax(out, dim=1)
probability = out[0, out.argmax(1).item()].item()
# 定义类别名称
classes = ['phone', 'picture']
# 获取预测类别和概率
predicted_class_index = out.argmax(1).item()
predicted_class_name = classes[predicted_class_index]
formatted_probability = "{:.2f}%".format(probability * 100) # 将概率转换为百分比格式
# 将预测结果转换为JSON格式
prediction = {
"class": predicted_class_name,
"probability": formatted_probability
}
prediction_json = json.dumps(prediction)
# 在控制台打印JSON格式的预测结果
print(prediction_json)
在第9行和第23行分别修改图片路径和训练好的模型路径，直接运行即可。预测结果下图为例，种类为手机UI，准确率100%。
总结
以上就是vgg16所复现的代码，本文仅仅简单对大佬的代码进行复现，非常感谢大佬的分享。如有纰漏，感谢批评指正
PS：增加类别数量要修改的位置：
txt.py第4行；
net.py第69行；
main.py第33行；
predict.py第27行
原文链接：[1]https://arxiv.org/abs/1409.1556
扫一扫
4766
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
