VGG16网络介绍及代码撰写详解(总结1)

[复制链接]
发表于 2024-10-1 20:20:57 | 显示全部楼层 |阅读模式

马上注册,结交更多好友,享用更多功能,让你轻松玩转社区。

您需要 登录 才可以下载或查看,没有账号?立即注册

×
  可以从本人以前的文章中可以看出作者以前从事的是嵌入式控制方面相干的工作,是一个机器视觉小白,之以是开始入门机器视觉的学习只要是一个idea,想把机器视觉与控制相融合将来做一点小东西。废话不多说开始正题。

摘要:本文是介绍VGG16网络,个人对其的知识总结,网络设计的知识点,以及代码怎样撰写,基于pytorch编写代码。作为一个刚入门的小白怎么去学习别人的代码,一步一步的去明白每一行代码,怎么将网络设计变成代码,模仿大佬的代码去撰写。作为小白如有不足之处请品评指正哈。
机器视觉基础知识

在此之前,本人学习了机器视觉的基础知识,以下是本人学习时的链接希望对你有所资助。
https://fuhanghang.blog.csdn.net/article/details/135544761?fromshare=blogdetail&sharetype=blogdetail&sharerId=135544761&sharerefer=PC&sharesource=weixin_52531699&sharefrom=from_link
【盘算机视觉与深度学习 北京邮电大学 鲁鹏 清楚版合集(完备版)】 https://www.bilibili.com/video/BV1V54y1B7K3/?share_source=copy_web&vd_source=b25ae79b699fbc0a2f70ccb983f6b74a
VGG16

在网络设计之前需要明白什么是VGG16。
以下是我鉴戒的文章的参考链接:
https://blog.csdn.net/weixin_46676835/article/details/129582927?fromshare=blogdetail&sharetype=blogdetail&sharerId=129582927&sharerefer=PC&sharesource=weixin_52531699&sharefrom=from_link
https://blog.csdn.net/weixin_46676835/article/details/128730174?fromshare=blogdetail&sharetype=blogdetail&sharerId=128730174&sharerefer=PC&sharesource=weixin_52531699&sharefrom=from_link

VGG16是一个深度卷积神经网络,它在2014年由牛津大学视觉多少组(Visual Geometry Group)提出,并在ImageNet图像分类使掷中取得了显著的结果。以下是VGG16的一些关键特点:
1. 网络结构
层数: VGG16包罗16个主要的权重层,包罗13个卷积层和3个全连接层。
卷积层: VGG16使用小的3x3卷积核进行卷积操作,增长了网络的深度,同时保持了较少的参数数目。
池化层: 每隔几个卷积层后,会使用2x2的最大池化层(Max Pooling)来降低特征图的尺寸,并减少盘算量。
2. 激活函数
VGG网络普遍使用ReLU(Rectified Linear Unit)作为激活函数,增长了网络的非线性表达能力。
3. 输入和输出
输入: VGG16接受224x224像素的RGB图像作为输入。
输出: 最终的输出是一个1000维的向量,表现属于1000个种别的概率分布(针对ImageNet数据集)。
4. 特点与长处
深度: VGG16相对较深(16层),使其能够学习更复杂的特征。
简单性: 使用同一的小卷积核和堆叠结构,使得网络设计相对简单且易于明白。
预训练模子: VGG16在多种使掷中被广泛应用,并且可以使用预训练权重进行迁徙学习,以提高其他使命的性能。
5. 应用方面
VGG16被广泛应用于图像分类、物体检测、图像分割等盘算机视觉使命。由于其精良的表现,很多后续的研究和模子(例如,VGG19等)都是在此基础上进行改进的。

从图中可以看出,VGG16包罗16个主要的权重层,包罗13个卷积层和3个全连接层。
盘算过程

1)输入图像尺寸为224x224x3,经64个通道为3的3x3的卷积核,步长为1,padding=same填充,卷积两次,再经ReLU激活,输出的尺寸大小为224x224x64
2)经max pooling(最大化池化),滤波器为2x2,步长为2,图像尺寸减半,池化后的尺寸变为112x112x64
3)经128个3x3的卷积核,两次卷积,ReLU激活,尺寸变为112x112x128
4)max pooling池化,尺寸变为56x56x128
5)经256个3x3的卷积核,三次卷积,ReLU激活,尺寸变为56x56x256
6)max pooling池化,尺寸变为28x28x256
7)经512个3x3的卷积核,三次卷积,ReLU激活,尺寸变为28x28x512
8)max pooling池化,尺寸变为14x14x512
9)经512个3x3的卷积核,三次卷积,ReLU,尺寸变为14x14x512
10)max pooling池化,尺寸变为7x7x512
11)然后Flatten(),将数据拉平成向量,变成一维512 ×7 ×7=25088。
11)再经过两层1x1x4096,一层1x1x1000的全连接层(共三层),经ReLU激活
12)末了通过softmax输出1000个预测效果,这是官方的(最终本人输出3种)
软件代码构思

有了以上理论基础后,开始构建代码思路,整体构建思路如下图所示,写代码之前肯定要构思好大抵思路,代码永远是为你思路框架服务的。

训练部分代码撰写

1.参数初始化
这一部分的代码的作用是,起首是关键参数以及头文件的初始化,再然后就是读取jpg文件,以及txt文件的内容,将其分为3类,蚂蚁,蜜蜂,狗,后续品种可以本身添加,jpg,txt训练数据多少都可以由本身添加,但是文件的命名格式必须为1.jpg,1.txt,2.jpg,2.txt才能读取文件。
(1)头文件初始化
  1. import os
  2. import parser
  3. import time
  4. import numpy as np
  5. import torch
  6. import torchvision.transforms as transforms
  7. import PIL
  8. from PIL import Image
  9. from scipy.stats import norm
  10. from torch import nn
  11. from torch.utils.data import Dataset, random_split, DataLoader, ConcatDataset
  12. import torchvision.datasets
  13. from torch import nn
  14. from torch.nn import MaxPool2d, Flatten
  15. import argparse
  16. from pathlib import Path
复制代码
import os:用于读取文件路径。
import time:用于时间相干的功能,本文用其盘算训练时间。
import torchvision.transforms as transforms:用于图像处置惩罚和转换的模块,常用于数据预处置惩罚。
import torchvision.datasets:提供了常用数据集的接口,例如MNIST、CIFAR等。
from PIL import Image:Pillow库,用于图像的打开、处置惩罚和保存。
from torch.nn import MaxPool2d, Flatten:导入特定的层,MaxPool2d用于进行二维最大池化,Flatten将多维输入展平为一维。
from torch.utils.data import Dataset, random_split, DataLoader, ConcatDataset:用于自界说数据集和数据加载器,方便批量处置惩罚数据。
import argparse:用于处置惩罚命令行参数,使得脚本可以接受用户输入的参数。
from pathlib import Path:提供高级路径操作功能,方便文件系统的操作。
(2)参数界说函数
用parse_opt()这个函数主要是见过很多大佬的文章都是用这个函数转义。作为小白像较为官方的标准化参数结构学习。
  1. def parse_opt():
  2.     parser = argparse.ArgumentParser()  # 创建 ArgumentParser 对象
  3.     parser.add_argument('--epochs', type=int, default=5, help='total training epochs')  # 添加参数
  4.     parser.add_argument('--batch_size', type=int, default=16, help='size of each batch')  # 添加批次大小参数
  5.     parser.add_argument('--learning_rate', type=int, default=0.03, help='size of learning_rate')  # 添加批次大小参数
  6.     #--device "cuda:0,cuda:1" 启用多个设备
  7.     parser.add_argument('--device', default='cuda:0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')  #--device cuda:1
  8.     # 解析参数
  9.     opt = parser.parse_args()
  10.     return opt
  11. opt = parse_opt()  # 调用解析函数
  12. epochs = opt.epochs              # 训练的轮数
  13. batch_size = opt.batch_size      # 每个批次的样本数量
  14. learning_rate = opt.learning_rate
  15. device = torch.device(opt.device)
复制代码
(3)界说文件处置惩罚类
这部分代码的编写主要是要实现一个功能,怎样把1.jpg,1.txt,2.jpg,2.txt…文件读取处置惩罚,并将txt文件的内容进行处置惩罚实现分类。
  1. class CustomDataset(Dataset):
  2.     def __init__(self, img_dir, label_dir, transform=None):
  3.         self.img_dir = img_dir
  4.         self.label_dir = label_dir
  5.         self.transform = transform
  6.         self.img_labels = self.load_labels()
  7.     def load_labels(self):
  8.         label_map = {'ants': 0, 'bees': 1,'dogs': 2}    #字典label_map 来定义标签与数字之间的映射关系
  9.         labels = []                                     #创建一个空列表 labels,用来存储从文件中读取到的标签的数字编码
  10.         for label_file in os.listdir(self.label_dir):
  11.             with open(os.path.join(self.label_dir, label_file), 'r') as f:  #对于每个标签文件,使用 open() 打开文件并读取第一行内容readline()
  12.                 line = f.readline().strip()         #strip() 方法用于去掉行首和行尾的空白字符(包括换行符)
  13.                 #从label_map 中获取当前行的标签对应的数字编码。如果当前行的内容不在 label_map 中,则返回-1
  14.                 labels.append(label_map.get(line, -1))
  15.         return labels
  16.     def __len__(self):
  17.         return len(self.img_labels)
  18.     def __getitem__(self, idx):
  19.         #使用 os.path.join 将图像目录路径 (self.img_dir) 和图像文件名拼接起来。文件名格式为 idx + 1(即从1开始计数),加上 .jpg 后缀。
  20.         #例如,如果idx为0,则img_path会是 D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_image\\1.jpg
  21.         img_path = os.path.join(self.img_dir, str(idx+1) + '.jpg')
  22.         image = Image.open(img_path).convert('RGB')  #打开指定路径的图像文件,并将其转换为 RGB 模式。
  23.         label = self.img_labels[idx]   #列表中获取当前图像的标签,idx 是当前图像的索引
  24.         if self.transform:   #transform变量是在CustomDataset中最终调用的transforms.Compose实现
  25.             image = self.transform(image)
  26.         return image, label
复制代码
使用字典 label_map 来界说标签与数字之间的映射关系:‘ants’ 对应 0;‘bees’ 对应 1;‘dogs’ 对应 2。对于每个标签文件,使用 open() 打开文件并读取第一行内容(readline())。strip() 方法用于去掉行首和行尾的空白字符(包罗换行符)。使用 label_map.get(line, -1) 从 label_map 中获取当前行的标签对应的数字编码。如果当前行的内容不在 label_map 中,则返回 -1。本人其时在这就犯了一个错误。可以从调试的框中看出,我有一个txt文件的内容并不是dogs致使其返回的值为-1,出现这个问题的缘故原由是本人txt文件中dogs没有保存,以是切记肯定要保存文件如果本身训练模子自建数据集的话。很多内容方面,解释写的很清楚,作者就不细讲了。

订正错误后的效果:

(4)数据集图片输入
代码需要实现一个图像大小设置功能,由于输入图像大小不确定需要将其设置成224×224大小图片,满足输入网络图像大小的设定。
  1. # 定义转换,包括调整大小
  2. transform = transforms.Compose([
  3.     transforms.Resize((224, 224)),  # 调整为 224x224 大小
  4.     transforms.ToTensor(),          # 转换为 Tensor
  5.     torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  6.     # 可以添加其他转换,例如归一化等
  7. ])
  8. # 初始化数据集
  9. ants_data = CustomDataset(
  10.     img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_image',
  11.     label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_label',
  12.     transform=transform
  13. )
  14. bees_data = CustomDataset(
  15.     img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\bees_image',
  16.     label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\bees_label',
  17.     transform=transform
  18. )
  19. dogs_data = CustomDataset(
  20.     img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\dogs_image',
  21.     label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\dogs_label',
  22.     transform=transform
  23. )
复制代码
transform 是通过 transforms.Compose 创建的一个组合转换对象。这个对象整合了多个图像预处置惩罚步骤,包罗调整大小、转换为张量,以及归一化。调整大小: 将图像调整为 224x224 像素。转换为 Tensor: 将 PIL 图像或 NumPy 数组转换为 PyTorch 张量。归一化: 根据给定的均值和标准差对张量进行归一化处置惩罚,使得数据分布更得当神经网络训练。归一化是VGG16网络的要求,但是本人发现没有也行,但所提供的代码为按照官方要求编写的。
ants_data 是 CustomDataset 类的一个实例,用于加载蚂蚁图像数据集。
img_dir 指定了存放图像文件的目次。
label_dir 指定了存放标签文件的目次。
transform 是应用于图像的转换操作(如缩放、裁剪、归一化等),即在加载图像时对其进行预处置惩罚。也就是调用transforms.Compose 。
从代码中也可以看出,实现的效果是读取图像读取标签然后通过transform将其变成224x224。

训练图片标签所在如图所示。
2.将数据集总和并随机打乱并分成训练集测试集
这段代码的作用是将我所读取的数据ants_data+bees_data+dogs_data总和起来,然后随机打乱,分割出训练集测试集分别是7:3。
  1. # 合并数据集
  2. total_data = ants_data+bees_data
  3. total_data = total_data+dogs_data
  4. # 计算训练集和验证集的大小
  5. total_size = len(total_data)
  6. train_size = int(0.7 * total_size)  # 70%
  7. val_size = total_size - train_size    # 30%
  8. train_data_size = train_size
  9. val_data_size = val_size
  10. print("训练集长度:{}".format(train_data_size))
  11. print("测试集长度:{}".format(val_data_size))
  12. # 随机分割数据集
  13. train_data, val_data = random_split(total_data, [train_size, val_size])
  14. # 创建 DataLoader,分成小批量(batches),以便于进行训练和验证
  15. train_dataloader = DataLoader(train_data, batch_size, shuffle=True) #shuffle=True可以随机打乱数据
  16. val_dataloader = DataLoader(val_data, batch_size, shuffle=False)
  17. # 打印数据集大小在这里插入图片描述
  18. print("训练集大小:{}".format(len(train_data)))
  19. print("验证集大小:{}".format(len(val_data)))
复制代码

从调试图中可知ants_data = 110,bees_data = 70,dogs_data = 120,total_data = 300,这也是我训练时提供的图片数目,标太多太累了。
3.创建网络模子
这部分是根据上图的vgg16网络结构搭建的。具体代码如下,本人写法为小白写法应该普通易懂。
  1. class VGG16(nn.Module):
  2.     def __init__(self):
  3.         super(VGG16, self).__init__()
  4.         self.model = nn.Sequential(  # 构建子类,可以减少foward代码撰写
  5.             nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
  6.             nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
  7.             nn.MaxPool2d(kernel_size=2, stride=2), #最大池化将224变成112
  8.             nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
  9.             nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
  10.             nn.MaxPool2d(kernel_size=2, stride=2), # 56x56
  11.             nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
  12.             nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
  13.             nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
  14.             nn.MaxPool2d(kernel_size=2, stride=2),  # 28x28
  15.             nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
  16.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  17.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  18.             nn.MaxPool2d(kernel_size=2, stride=2), # 14x14
  19.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  20.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  21.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  22.             nn.MaxPool2d(kernel_size=2, stride=2),  # 7x7
  23.         )
  24.         # 使用AdaptiveAvgPool2d来确保输出为固定大小
  25.         self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
  26.         # 全连接层
  27.         self.fc1 = nn.Linear(in_features=512 * 7 * 7, out_features=4096)
  28.         self.fc2 = nn.Linear(in_features=4096, out_features=4096)
  29.         self.fc3 = nn.Linear(in_features=4096, out_features=3)
  30.         self.relu = nn.ReLU()
  31.         self.dropout = nn.Dropout(0.3)
  32.     def forward(self, x):
  33.         x = self.model(x)
  34.         x = self.avgpool(x)  # 使用自适应平均池化
  35.         x = x.view(x.size(0), -1)  # 拉平操作,(batch_size, num_features) 也就是(8,25088) 512*7*7 = 25088
  36.         x = self.relu(self.fc1(x))
  37.         x = self.dropout(x)
  38.         x = self.relu(self.fc2(x))
  39.         x = self.dropout(x)
  40.         x = self.fc3(x)
  41.         return x
复制代码
  1. x = self.avgpool(x) 自适应平均池化,将输入特征图(x)的空间维度降低到固定大小。
复制代码
  1. x = x.view(x.size(0), -1)  这行代码将输出张量 x 的形状变为 (batch_size, num_features),
  2. 其中 batch_size 是当前批次的样本数量,num_features 是特征的总数。这一步是为了将多维张量转换为二维张量,
  3. 以便进入全连接层进行分类。
复制代码
4.训练加测试
这部分得先明白工作步骤再写代码。1.训练集训练调用VGG16网络结构,丧失函数构建,随机梯度下降,优化训练 2.测试集测试,设置为评估模式,调用VGG16网络结构,丧失函数构建,丧失盘算。3.输出训练丧失,测试丧失,正确率。这部分是机器视觉基础概念知识。
  1. def train():
  2.     My_VGG16 = VGG16()              #调用类
  3.     My_VGG16 = My_VGG16.to(device)
  4.     # (1).损失函数构建
  5.     loss_fn = nn.CrossEntropyLoss()  #计算预测值与真实标签之间的差异
  6.     loss_fn = loss_fn.to(device)     #将模型和数据都放在同一个设备上,GPU
  7.     # (2).优化器
  8.     # #随机梯度下降(Stochastic Gradient Descent)优化器的一种实现。SGD 是一种常见的优化算法
  9.     optimizer = torch.optim.SGD(My_VGG16.parameters(), lr=learning_rate)
  10.     # 用于存储损失和准确率
  11.     train_losses = []
  12.     val_losses = []
  13.     accuracies = []
  14.     for i in range(epochs):
  15.         loss_temp = 0  # 临时变量
  16.         print("--------第{}轮训练开始--------".format(i + 1))
  17.         # 训练阶段
  18.         My_VGG16.train()  # 设置为训练模式
  19.         #data有两个部分分别是(8,3,224,224),(8,) 一定要注意顺序
  20.         for data in train_dataloader:
  21.             imgs, targets = data
  22.             imgs = imgs.to(device)
  23.             targets = targets.to(device)
  24.             #outputs = tensor(8,3)
  25.             outputs = My_VGG16(imgs)
  26.             loss = loss_fn(outputs, targets)
  27.             # 优化器优化模型
  28.             optimizer.zero_grad()  # 梯度清零
  29.             loss.backward()  # 反向传播
  30.             optimizer.step()  # 梯度更新
  31.             loss_temp += loss.item()
  32.         # 记录训练损失
  33.         train_losses.append(loss_temp / len(train_dataloader))
  34.         # 测试阶段
  35.         My_VGG16.eval()  # 设置为评估模式
  36.         total_test_loss = 0
  37.         total_accuracy = 0
  38.         with torch.no_grad():
  39.             for data in val_dataloader:
  40.                 imgs, targets = data
  41.                 imgs = imgs.to(device)
  42.                 targets = targets.to(device)
  43.                 outputs = My_VGG16(imgs)
  44.                 loss = loss_fn(outputs, targets)
  45.                 total_test_loss += loss.item()
  46.                 #找到每个样本的预测类别,然后与真实标签进行比较
  47.                 accuracy = (outputs.argmax(1) == targets).sum().item()  # 使用.item()将Tensor转换为Python数值
  48.                 total_accuracy += accuracy   #计算正确预测的数量并累加到 total_accuracy
复制代码
  1. imgs, targets = data这段代码
复制代码

可以从图中看出data有两个部分分别是(8,3,224,224),(8,) 肯定要留意序次分别对应的就是imgs,targets
  1. outputs = My_VGG16(imgs)  outputs = tensor(8,3)
复制代码

outputs单纯的就是输出品种效果,8是因为我的batch_size = 8,输入的是8张图片。
5.保存数据至.pth文件并画图
这部分代码的编写需要实现的是,将训练的权值保存至.pth文件中,并在窗口中显示训练时的时间,loss,正确率输出思量到train_loss为list:5,以是直接使用循环的打印输出效果。

  1.   # 记录验证损失和准确率
  2.         val_losses.append(total_test_loss / len(val_dataloader))
  3.         accuracies.append(total_accuracy / val_data_size)
  4.         print("整体测试集上的正确率:{}".format(total_accuracy / val_data_size))
  5.         print("整体测试集上的Loss:{}".format(total_test_loss))
  6.     # 保存模型
  7.     torch.save(My_VGG16, "vgg16_{}.pth".format(2))
  8.     torch.save(My_VGG16.state_dict(), "vgg16_dict_{}.txt".format(2))
  9.     print("模型已保存")
  10.     # 绘制损失和准确率曲线
  11.     epochs_range = range(1, epochs + 1)
  12.     plt.figure(figsize=(12, 5))  #创建一个新的图形窗口,设置图形的大小
  13.     # 训练和验证损失
  14.     plt.subplot(1, 2, 1)  #第一个图像框
  15.     plt.plot(epochs_range, train_losses, label='Training Loss')
  16.     #train_losses=[1.0952595781396937, 1.085733965591148, 1.07982603708903, 1.0776626953372248, 1.0816458507820412]
  17.     plt.plot(epochs_range, val_losses, label='Validation Loss')
  18.     plt.title('Loss vs Epochs')
  19.     plt.xlabel('Epochs')
  20.     plt.ylabel('Loss')
  21.     plt.legend()
  22.     # 准确率
  23.     plt.subplot(1, 2, 2)
  24.     plt.plot(epochs_range, accuracies, label='Validation Accuracy')
  25.     plt.title('Accuracy vs Epochs')
  26.     plt.xlabel('Epochs')
  27.     plt.ylabel('Accuracy')
  28.     plt.legend()
  29.     plt.tight_layout()
  30.     plt.show()
  31. if __name__ == '__main__':
  32.     parse_opt()
  33.     train()
  34.     # 测试模型
  35.     # model = VGG16()
  36.     # sample_input = torch.randn(1, 3, 224, 224)  # Batch size 为 1
  37.     # output = model(sample_input)
  38.     # print("Final output shape:", output.shape)  # 应输出 [1, 3]
复制代码
代码讲解,解释很具体,不做细数。

输出图片效果,这是只训练了5次的网络,可以看到训练效果是不理想的。
以上就是训练集代码及讲解。
训练整体代码:
  1. #VGG_16训练网络模子,第一代测试代码,主要功能如下#辨认3类蚂蚁,蜜蜂,狗,可恣意添加,文件名为1.jpg,1.txt,内容ants#2代是在一代的基础上添加了图像显示,标准化参数设计方面import os
  2. import parser
  3. import time
  4. import numpy as np
  5. import torch
  6. import torchvision.transforms as transforms
  7. import PIL
  8. from PIL import Image
  9. from scipy.stats import norm
  10. from torch import nn
  11. from torch.utils.data import Dataset, random_split, DataLoader, ConcatDataset
  12. import torchvision.datasets
  13. from torch import nn
  14. from torch.nn import MaxPool2d, Flatten
  15. import argparse
  16. from pathlib import Path
  17. # ------------------------------1.初始化,标准参数初始化-------------------------------def parse_opt():    parser = argparse.ArgumentParser()  # 创建 ArgumentParser 对象    parser.add_argument('--epochs', type=int, default=5, help='total training epochs')  # 添加参数    parser.add_argument('--batch_size', type=int, default=8, help='size of each batch')  # 添加批次大小参数    parser.add_argument('--learning_rate', type=int, default=0.03, help='size of learning_rate')  # 添加批次大小参数    #--device "cuda:0,cuda:1" 启用多个设备    parser.add_argument('--device', default='cuda:0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')  #--device cuda:1    # 剖析参数    opt = parser.parse_args()    return optopt = parse_opt()  # 调用剖析函数epochs = opt.epochs              # 训练的轮数batch_size = opt.batch_size      # 每个批次的样本数目learning_rate = opt.learning_ratedevice = torch.device(opt.device)class CustomDataset(Dataset):
  18.     def __init__(self, img_dir, label_dir, transform=None):
  19.         self.img_dir = img_dir
  20.         self.label_dir = label_dir
  21.         self.transform = transform
  22.         self.img_labels = self.load_labels()
  23.     def load_labels(self):
  24.         label_map = {'ants': 0, 'bees': 1,'dogs': 2}    #字典label_map 来定义标签与数字之间的映射关系
  25.         labels = []                                     #创建一个空列表 labels,用来存储从文件中读取到的标签的数字编码
  26.         for label_file in os.listdir(self.label_dir):
  27.             with open(os.path.join(self.label_dir, label_file), 'r') as f:  #对于每个标签文件,使用 open() 打开文件并读取第一行内容readline()
  28.                 line = f.readline().strip()         #strip() 方法用于去掉行首和行尾的空白字符(包括换行符)
  29.                 #从label_map 中获取当前行的标签对应的数字编码。如果当前行的内容不在 label_map 中,则返回-1
  30.                 labels.append(label_map.get(line, -1))
  31.         return labels
  32.     def __len__(self):
  33.         return len(self.img_labels)
  34.     def __getitem__(self, idx):
  35.         #使用 os.path.join 将图像目录路径 (self.img_dir) 和图像文件名拼接起来。文件名格式为 idx + 1(即从1开始计数),加上 .jpg 后缀。
  36.         #例如,如果idx为0,则img_path会是 D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_image\\1.jpg
  37.         img_path = os.path.join(self.img_dir, str(idx+1) + '.jpg')
  38.         image = Image.open(img_path).convert('RGB')  #打开指定路径的图像文件,并将其转换为 RGB 模式。
  39.         label = self.img_labels[idx]   #列表中获取当前图像的标签,idx 是当前图像的索引
  40.         if self.transform:   #transform变量是在CustomDataset中最终调用的transforms.Compose实现
  41.             image = self.transform(image)
  42.         return image, label
  43. # 定义转换,包括调整大小
  44. transform = transforms.Compose([
  45.     transforms.Resize((224, 224)),  # 调整为 224x224 大小
  46.     transforms.ToTensor(),          # 转换为 Tensor
  47.     torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  48.     # 可以添加其他转换,例如归一化等
  49. ])
  50. # 初始化数据集
  51. ants_data = CustomDataset(
  52.     img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_image',
  53.     label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_label',
  54.     transform=transform
  55. )
  56. bees_data = CustomDataset(
  57.     img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\bees_image',
  58.     label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\bees_label',
  59.     transform=transform
  60. )
  61. dogs_data = CustomDataset(
  62.     img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\dogs_image',
  63.     label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\dogs_label',
  64.     transform=transform
  65. )
  66. #-----------------------2.将数据集总和并随机打乱并分成训练集测试集-------------------# 归并数据集total_data = ants_data+bees_datatotal_data = total_data+dogs_data# 盘算训练集和验证集的大小total_size = len(total_data)train_size = int(0.7 * total_size)  # 70%val_size = total_size - train_size    # 30%train_data_size = train_sizeval_data_size = val_sizeprint("训练集长度:{}".format(train_data_size))print("测试集长度:{}".format(val_data_size))# 随机分割数据集train_data, val_data = random_split(total_data, [train_size, val_size])# 创建 DataLoader,分成小批量(batches),以便于进行训练和验证train_dataloader = DataLoader(train_data, batch_size, shuffle=True) #shuffle=True可以随机打乱数据val_dataloader = DataLoader(val_data, batch_size, shuffle=False)# 打印数据集大小print("训练集大小:{}".format(len(train_data)))print("验证集大小:{}".format(len(val_data)))#----------------------------3.创建网络模子--------------------------------class VGG16(nn.Module):
  67.     def __init__(self):
  68.         super(VGG16, self).__init__()
  69.         self.model = nn.Sequential(  # 构建子类,可以减少foward代码撰写
  70.             nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
  71.             nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
  72.             nn.MaxPool2d(kernel_size=2, stride=2), #最大池化将224变成112
  73.             nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
  74.             nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
  75.             nn.MaxPool2d(kernel_size=2, stride=2), # 56x56
  76.             nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
  77.             nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
  78.             nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
  79.             nn.MaxPool2d(kernel_size=2, stride=2),  # 28x28
  80.             nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
  81.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  82.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  83.             nn.MaxPool2d(kernel_size=2, stride=2), # 14x14
  84.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  85.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  86.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  87.             nn.MaxPool2d(kernel_size=2, stride=2),  # 7x7
  88.         )
  89.         # 使用AdaptiveAvgPool2d来确保输出为固定大小
  90.         self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
  91.         # 全连接层
  92.         self.fc1 = nn.Linear(in_features=512 * 7 * 7, out_features=4096)
  93.         self.fc2 = nn.Linear(in_features=4096, out_features=4096)
  94.         self.fc3 = nn.Linear(in_features=4096, out_features=3)
  95.         self.relu = nn.ReLU()
  96.         self.dropout = nn.Dropout(0.3)
  97.     def forward(self, x):
  98.         x = self.model(x)
  99.         x = self.avgpool(x)  # 使用自适应平均池化
  100.         x = x.view(x.size(0), -1)  # 拉平操作,(batch_size, num_features) 也就是(8,25088) 512*7*7 = 25088
  101.         x = self.relu(self.fc1(x))
  102.         x = self.dropout(x)
  103.         x = self.relu(self.fc2(x))
  104.         x = self.dropout(x)
  105.         x = self.fc3(x)
  106.         return x
  107. #------------------------------------4.训练加测试-----------------------import numpy as npimport matplotlibmatplotlib.use('TkAgg')  # 大概尝试 'Qt5Agg',有这行代码会多一个弹窗显示import matplotlib.pyplot as pltimport torch.nn as nndef train():
  108.     My_VGG16 = VGG16()              #调用类
  109.     My_VGG16 = My_VGG16.to(device)
  110.     # (1).损失函数构建
  111.     loss_fn = nn.CrossEntropyLoss()  #计算预测值与真实标签之间的差异
  112.     loss_fn = loss_fn.to(device)     #将模型和数据都放在同一个设备上,GPU
  113.     # (2).优化器
  114.     # #随机梯度下降(Stochastic Gradient Descent)优化器的一种实现。SGD 是一种常见的优化算法
  115.     optimizer = torch.optim.SGD(My_VGG16.parameters(), lr=learning_rate)
  116.     # 用于存储损失和准确率
  117.     train_losses = []
  118.     val_losses = []
  119.     accuracies = []
  120.     for i in range(epochs):
  121.         loss_temp = 0  # 临时变量
  122.         print("--------第{}轮训练开始--------".format(i + 1))
  123.         # 训练阶段
  124.         My_VGG16.train()  # 设置为训练模式
  125.         #data有两个部分分别是(8,3,224,224),(8,) 一定要注意顺序
  126.         for data in train_dataloader:
  127.             imgs, targets = data
  128.             imgs = imgs.to(device)
  129.             targets = targets.to(device)
  130.             #outputs = tensor(8,3)
  131.             outputs = My_VGG16(imgs)
  132.             loss = loss_fn(outputs, targets)
  133.             # 优化器优化模型
  134.             optimizer.zero_grad()  # 梯度清零
  135.             loss.backward()  # 反向传播
  136.             optimizer.step()  # 梯度更新
  137.             loss_temp += loss.item()
  138.         # 记录训练损失
  139.         train_losses.append(loss_temp / len(train_dataloader))
  140.         # 测试阶段
  141.         My_VGG16.eval()  # 设置为评估模式
  142.         total_test_loss = 0
  143.         total_accuracy = 0
  144.         with torch.no_grad():
  145.             for data in val_dataloader:
  146.                 imgs, targets = data
  147.                 imgs = imgs.to(device)
  148.                 targets = targets.to(device)
  149.                 outputs = My_VGG16(imgs)
  150.                 loss = loss_fn(outputs, targets)
  151.                 total_test_loss += loss.item()
  152.                 #找到每个样本的预测类别,然后与真实标签进行比较
  153.                 accuracy = (outputs.argmax(1) == targets).sum().item()  # 使用.item()将Tensor转换为Python数值
  154.                 total_accuracy += accuracy   #计算正确预测的数量并累加到 total_accuracy
  155. #-----------------------5.保存数据至.pth文件并画图---------------------------------------------------        # 记录验证损失和准确率
  156.         val_losses.append(total_test_loss / len(val_dataloader))
  157.         accuracies.append(total_accuracy / val_data_size)
  158.         print("整体测试集上的正确率:{}".format(total_accuracy / val_data_size))
  159.         print("整体测试集上的Loss:{}".format(total_test_loss))
  160.     # 保存模型
  161.     torch.save(My_VGG16, "vgg16_{}.pth".format(2))
  162.     torch.save(My_VGG16.state_dict(), "vgg16_dict_{}.txt".format(2))
  163.     print("模型已保存")
  164.     # 绘制损失和准确率曲线
  165.     epochs_range = range(1, epochs + 1)
  166.     plt.figure(figsize=(12, 5))  #创建一个新的图形窗口,设置图形的大小
  167.     # 训练和验证损失
  168.     plt.subplot(1, 2, 1)  #第一个图像框
  169.     plt.plot(epochs_range, train_losses, label='Training Loss')
  170.     #train_losses=[1.0952595781396937, 1.085733965591148, 1.07982603708903, 1.0776626953372248, 1.0816458507820412]
  171.     plt.plot(epochs_range, val_losses, label='Validation Loss')
  172.     plt.title('Loss vs Epochs')
  173.     plt.xlabel('Epochs')
  174.     plt.ylabel('Loss')
  175.     plt.legend()
  176.     # 准确率
  177.     plt.subplot(1, 2, 2)
  178.     plt.plot(epochs_range, accuracies, label='Validation Accuracy')
  179.     plt.title('Accuracy vs Epochs')
  180.     plt.xlabel('Epochs')
  181.     plt.ylabel('Accuracy')
  182.     plt.legend()
  183.     plt.tight_layout()
  184.     plt.show()
  185. if __name__ == '__main__':
  186.     parse_opt()
  187.     train()
  188.     # 测试模型
  189.     # model = VGG16()
  190.     # sample_input = torch.randn(1, 3, 224, 224)  # Batch size 为 1
  191.     # output = model(sample_input)
  192.     # print("Final output shape:", output.shape)  # 应输出 [1, 3]
复制代码
留意本人的训练是用的gpu,不推荐用cpu跑代码训练。
检测部分代码撰写

这一部分的代码主要实现的是输入恣意大小的图片,预测出它的种类。
1.导入头文件参数初始化
导入头文件,参数标准化后,读取图片,并通过transform实现图片输入格式的设置。

  1. import torch
  2. import torchvision
  3. from PIL import Image
  4. from torch import nn
  5. from torch.nn import MaxPool2d
  6. import argparse
  7. from pathlib import Path
  8. #--------------------------1.参数数据初始化-----------------------------------
  9. def parse_opt():
  10.     parser = argparse.ArgumentParser()  # 创建 ArgumentParser 对象
  11.     parser.add_argument('--Dropout_Val', type=int, default=0.4, help='Dropout_Val number')  # 添加参数
  12.     #--device "cuda:0,cuda:1" 启用多个设备
  13.     parser.add_argument('--device', default='cuda:0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')  #--device cuda:1
  14.     # 解析参数
  15.     opt = parser.parse_args()
  16.     return opt
  17. opt = parse_opt()  # 调用解析函数
  18. Dropout_Val = opt.Dropout_Val
  19. device = torch.device(opt.device)
  20. image_path = "D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\1.jpg"
  21. image = Image.open(image_path)   #注意格式为RGBA
  22. print(image)
  23. image = image.convert("RGB")     #将格式变成RGB
  24. print(image)
  25. #确保你的图像预处理步骤(如归一化)与模型训练时使用的步骤相匹配。通常,VGG16 需要将图像归一化到 [0, 1] 范围内,
  26. # 并可能需要减去均值和除以标准差
  27. transform = torchvision.transforms.Compose([
  28.     torchvision.transforms.Resize((224, 224)),
  29.     torchvision.transforms.ToTensor(),
  30.     torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  31. ])
  32. image = transform(image)
  33. print(image.shape)
  34. # 类别标签
  35. class_names = ['ant', 'bee', 'dog']  # 根据实际情况填写
复制代码
2.搭建VGG16网络结构
  1. class VGG16(nn.Module):
  2.     def __init__(self):
  3.         super(VGG16, self).__init__()
  4.         self.model = nn.Sequential(  # 构建子类,可以减少foward代码撰写
  5.             nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
  6.             nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
  7.             nn.MaxPool2d(kernel_size=2, stride=2), #最大池化将224变成112
  8.             nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
  9.             nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
  10.             nn.MaxPool2d(kernel_size=2, stride=2), # 56x56
  11.             nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
  12.             nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
  13.             nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
  14.             nn.MaxPool2d(kernel_size=2, stride=2),  # 28x28
  15.             nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
  16.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  17.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  18.             nn.MaxPool2d(kernel_size=2, stride=2), # 14x14
  19.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  20.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  21.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  22.             nn.MaxPool2d(kernel_size=2, stride=2),  # 7x7
  23.         )
  24.         # 使用AdaptiveAvgPool2d来确保输出为固定大小
  25.         self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
  26.         # 全连接层
  27.         self.fc1 = nn.Linear(in_features=512 * 7 * 7, out_features=4096)
  28.         self.fc2 = nn.Linear(in_features=4096, out_features=4096)
  29.         self.fc3 = nn.Linear(in_features=4096, out_features=3)
  30.         self.relu = nn.ReLU()
  31.         self.dropout = nn.Dropout(Dropout_Val)
  32.     def forward(self, x):
  33.         x = self.model(x)
  34.         x = self.avgpool(x)  # 使用自适应平均池化
  35.         x = x.view(x.size(0), -1)  # 拉平操作
  36.         x = self.relu(self.fc1(x))
  37.         x = self.dropout(x)
  38.         x = self.relu(self.fc2(x))
  39.         x = self.dropout(x)
  40.         x = self.fc3(x)
  41.         return x
复制代码
3.输出测试效果
输出测试效果需要将图片送入VGG16网络进行盘算,带入训练好的权值vgg16_1.pth,得出预测效果,将预测的效果输出成概率从而盘算是什么品种。
  1. model = VGG16()
  2. model = torch.load("vgg16_1.pth")
  3. #模型的权重在GPU上,而输入张量在CPU上。要解决这个问题,你可以将输入张量移动到GPU
  4. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  5. model.to(device)
  6. image = image.to(device)
  7. #模型的权重在CPU上,而输入张量在GPU上。要解决这个问题,你可以将输入张量移动到CPU
  8. #model = torch.load("tudui_0.pth",map_location=torch.device('cpu'))
  9. image = torch.reshape(image,(1,3,224,224))
  10. model.eval()
  11. with torch.no_grad():
  12.     output = model(image)
  13. predicted_index = output.argmax(1).item()  # 获取预测的索引
  14. print(output.argmax(1))
  15. print(output)   #输出结果预测的是第六个类别
  16. print("Predicted index:", predicted_index)
  17. # 输出对应的类别名称
  18. print("Predicted class:", class_names[predicted_index])
复制代码
测试输出效果



着实多训练频频预测效果还是可以的。
测试整体代码
  1. # vgg16网络测试文件,第一代测试代码,主要功能如下# 恣意输入一张图片,输出判断范例,二代与一代相比就是多了参数标准化
  2. import torch
  3. import torchvision
  4. from PIL import Image
  5. from torch import nn
  6. from torch.nn import MaxPool2d
  7. import argparse
  8. from pathlib import Path
  9. #--------------------------1.参数数据初始化-----------------------------------
  10. def parse_opt():
  11.     parser = argparse.ArgumentParser()  # 创建 ArgumentParser 对象
  12.     parser.add_argument('--Dropout_Val', type=int, default=0.4, help='Dropout_Val number')  # 添加参数
  13.     #--device "cuda:0,cuda:1" 启用多个设备
  14.     parser.add_argument('--device', default='cuda:0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')  #--device cuda:1
  15.     # 解析参数
  16.     opt = parser.parse_args()
  17.     return opt
  18. opt = parse_opt()  # 调用解析函数
  19. Dropout_Val = opt.Dropout_Val
  20. device = torch.device(opt.device)
  21. image_path = "D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\1.jpg"
  22. image = Image.open(image_path)   #注意格式为RGBA
  23. print(image)
  24. image = image.convert("RGB")     #将格式变成RGB
  25. print(image)
  26. #确保你的图像预处理步骤(如归一化)与模型训练时使用的步骤相匹配。通常,VGG16 需要将图像归一化到 [0, 1] 范围内,
  27. # 并可能需要减去均值和除以标准差
  28. transform = torchvision.transforms.Compose([
  29.     torchvision.transforms.Resize((224, 224)),
  30.     torchvision.transforms.ToTensor(),
  31.     torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  32. ])
  33. image = transform(image)
  34. print(image.shape)
  35. # 类别标签
  36. class_names = ['ant', 'bee', 'dog']  # 根据实际情况填写
  37. #-------------------------------------2.搭建VGG16网络结构---------------------------------------------class VGG16(nn.Module):
  38.     def __init__(self):
  39.         super(VGG16, self).__init__()
  40.         self.model = nn.Sequential(  # 构建子类,可以减少foward代码撰写
  41.             nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
  42.             nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
  43.             nn.MaxPool2d(kernel_size=2, stride=2), #最大池化将224变成112
  44.             nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
  45.             nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
  46.             nn.MaxPool2d(kernel_size=2, stride=2), # 56x56
  47.             nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
  48.             nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
  49.             nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
  50.             nn.MaxPool2d(kernel_size=2, stride=2),  # 28x28
  51.             nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
  52.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  53.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  54.             nn.MaxPool2d(kernel_size=2, stride=2), # 14x14
  55.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  56.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  57.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
  58.             nn.MaxPool2d(kernel_size=2, stride=2),  # 7x7
  59.         )
  60.         # 使用AdaptiveAvgPool2d来确保输出为固定大小
  61.         self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
  62.         # 全连接层
  63.         self.fc1 = nn.Linear(in_features=512 * 7 * 7, out_features=4096)
  64.         self.fc2 = nn.Linear(in_features=4096, out_features=4096)
  65.         self.fc3 = nn.Linear(in_features=4096, out_features=3)
  66.         self.relu = nn.ReLU()
  67.         self.dropout = nn.Dropout(Dropout_Val)
  68.     def forward(self, x):
  69.         x = self.model(x)
  70.         x = self.avgpool(x)  # 使用自适应平均池化
  71.         x = x.view(x.size(0), -1)  # 拉平操作
  72.         x = self.relu(self.fc1(x))
  73.         x = self.dropout(x)
  74.         x = self.relu(self.fc2(x))
  75.         x = self.dropout(x)
  76.         x = self.fc3(x)
  77.         return x
  78. #---------------------------3.输出测试效果-----------------------------model = VGG16()
  79. model = torch.load("vgg16_1.pth")
  80. #模型的权重在GPU上,而输入张量在CPU上。要解决这个问题,你可以将输入张量移动到GPU
  81. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  82. model.to(device)
  83. image = image.to(device)
  84. #模型的权重在CPU上,而输入张量在GPU上。要解决这个问题,你可以将输入张量移动到CPU
  85. #model = torch.load("tudui_0.pth",map_location=torch.device('cpu'))
  86. image = torch.reshape(image,(1,3,224,224))
  87. model.eval()
  88. with torch.no_grad():
  89.     output = model(image)
  90. predicted_index = output.argmax(1).item()  # 获取预测的索引
  91. print(output.argmax(1))
  92. print(output)   #输出结果预测的是第六个类别
  93. print("Predicted index:", predicted_index)
  94. # 输出对应的类别名称
  95. print("Predicted class:", class_names[predicted_index])
复制代码
最开始训练的次数有点少,导致预测效果不理想,厥后去用饭前设置训练epochs = 500次,返来发现效果变好了一点。提高正确率的方式很多,还可以丰富训练集。留意我的权值文件权值vgg16_1.pth,大概使用的是权值vgg16_2.pth,代码我就不改了。主要是要明白以上代码。
以上就是本人的心得与总结,如有不足之处请多多包涵。
百度网盘链接代码权值及训练图片: https://pan.baidu.com/s/14xixodOGVEMK-KpE5rZTtw
提取码: dwrg
回复

使用道具 举报

登录后关闭弹窗

登录参与点评抽奖  加入IT实名职场社区
去登录
快速回复 返回顶部 返回列表