VGG16网络介绍及代码撰写详解（总结1）

可以从本人以前的文章中可以看出作者以前从事的是嵌入式控制方面相关的工作，是一个机器视觉小白，之所以开始入门机器视觉的学习只要是一个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）头文件初始化

import os
import parser
import time
import numpy as np
import torch
import torchvision.transforms as transforms
import PIL
from PIL import Image
from scipy.stats import norm
from torch import nn
from torch.utils.data import Dataset, random_split, DataLoader, ConcatDataset
import torchvision.datasets
from torch import nn
from torch.nn import MaxPool2d, Flatten

import argparse
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()这个函数主要是见过许多大佬的文章都是用这个函数转义。作为小白像较为官方的标准化参数结构学习。

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=16, 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 opt

opt = parse_opt()  # 调用解析函数
epochs = opt.epochs              # 训练的轮数
batch_size = opt.batch_size      # 每个批次的样本数量
learning_rate = opt.learning_rate
device = torch.device(opt.device)

（3）定义文件处理类
这部分代码的编写主要是要实现一个功能，如何把1.jpg，1.txt，2.jpg，2.txt...文件读取处理，并将txt文件的内容进行处理实现分类。

class CustomDataset(Dataset):
    def __init__(self, img_dir, label_dir, transform=None):
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.transform = transform
        self.img_labels = self.load_labels()

    def load_labels(self):
        label_map = {'ants': 0, 'bees': 1,'dogs': 2}    #字典label_map 来定义标签与数字之间的映射关系
        labels = []                                     #创建一个空列表 labels，用来存储从文件中读取到的标签的数字编码
        for label_file in os.listdir(self.label_dir):
            with open(os.path.join(self.label_dir, label_file), 'r') as f:  #对于每个标签文件，使用 open() 打开文件并读取第一行内容readline()
                line = f.readline().strip()         #strip() 方法用于去掉行首和行尾的空白字符（包括换行符）
                #从label_map 中获取当前行的标签对应的数字编码。如果当前行的内容不在 label_map 中，则返回-1
                labels.append(label_map.get(line, -1))
        return labels

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        #使用 os.path.join 将图像目录路径 (self.img_dir) 和图像文件名拼接起来。文件名格式为 idx + 1（即从1开始计数），加上 .jpg 后缀。
        #例如，如果idx为0，则img_path会是 D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_image\\1.jpg
        img_path = os.path.join(self.img_dir, str(idx+1) + '.jpg')
        image = Image.open(img_path).convert('RGB')  #打开指定路径的图像文件，并将其转换为 RGB 模式。
        label = self.img_labels[idx]   #列表中获取当前图像的标签，idx 是当前图像的索引

        if self.transform:   #transform变量是在CustomDataset中最终调用的transforms.Compose实现
            image = self.transform(image)

        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大小图片，满足输入网络图像大小的设定。

# 定义转换，包括调整大小
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整为 224x224 大小
    transforms.ToTensor(),          # 转换为 Tensor
    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    # 可以添加其他转换，例如归一化等
])

# 初始化数据集
ants_data = CustomDataset(
    img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_image',
    label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_label',
    transform=transform
)


bees_data = CustomDataset(
    img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\bees_image',
    label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\bees_label',
    transform=transform
)

dogs_data = CustomDataset(
    img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\dogs_image',
    label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\dogs_label',
    transform=transform
)

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。

# 合并数据集
total_data = ants_data+bees_data
total_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_size
val_data_size = val_size

print("训练集长度：{}".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)))

从调试图中可知ants_data = 110,bees_data = 70,dogs_data = 120,total_data = 300，这也是我训练时提供的图片数目，标太多太累了。

3.创建网络模型
这部分是根据上图的vgg16网络结构搭建的。具体代码如下，本人写法为小白写法应该通俗易懂。

class VGG16(nn.Module):
    def __init__(self):
        super(VGG16, self).__init__()
        self.model = nn.Sequential(  # 构建子类，可以减少foward代码撰写
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), #最大池化将224变成112
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), # 56x56
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 28x28
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), # 14x14
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 7x7
        )
        # 使用AdaptiveAvgPool2d来确保输出为固定大小
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))

        # 全连接层
        self.fc1 = nn.Linear(in_features=512 * 7 * 7, out_features=4096)
        self.fc2 = nn.Linear(in_features=4096, out_features=4096)
        self.fc3 = nn.Linear(in_features=4096, out_features=3)

        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.3)

    def forward(self, x):
        x = self.model(x)
        x = self.avgpool(x)  # 使用自适应平均池化
        x = x.view(x.size(0), -1)  # 拉平操作,(batch_size, num_features) 也就是（8，25088） 512*7*7 = 25088
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x

x = self.avgpool(x) 自适应平均池化，将输入特征图（x）的空间维度降低到固定大小。

x = x.view(x.size(0), -1)  这行代码将输出张量 x 的形状变为 (batch_size, num_features)，
其中 batch_size 是当前批次的样本数量，num_features 是特征的总数。这一步是为了将多维张量转换为二维张量，
以便进入全连接层进行分类。

4.训练加测试
这部分得先明白工作步骤再写代码。1.训练集训练调用VGG16网络结构，损失函数构建，随机梯度下降，优化训练 2.测试集测试，设置为评估模式，调用VGG16网络结构，损失函数构建，损失计算。3.输出训练损失，测试损失，准确率。这部分是机器视觉基础概念知识。

def train():
    My_VGG16 = VGG16()              #调用类
    My_VGG16 = My_VGG16.to(device)

    # (1).损失函数构建
    loss_fn = nn.CrossEntropyLoss()  #计算预测值与真实标签之间的差异
    loss_fn = loss_fn.to(device)     #将模型和数据都放在同一个设备上，GPU

    # (2).优化器
    # #随机梯度下降（Stochastic Gradient Descent）优化器的一种实现。SGD 是一种常见的优化算法
    optimizer = torch.optim.SGD(My_VGG16.parameters(), lr=learning_rate)


    # 用于存储损失和准确率
    train_losses = []
    val_losses = []
    accuracies = []

    for i in range(epochs):
        loss_temp = 0  # 临时变量
        print("--------第{}轮训练开始--------".format(i + 1))

        # 训练阶段
        My_VGG16.train()  # 设置为训练模式
        #data有两个部分分别是(8,3,224,224),(8,) 一定要注意顺序
        for data in train_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)
            #outputs = tensor（8，3）
            outputs = My_VGG16(imgs)
            loss = loss_fn(outputs, targets)

            # 优化器优化模型
            optimizer.zero_grad()  # 梯度清零
            loss.backward()  # 反向传播
            optimizer.step()  # 梯度更新

            loss_temp += loss.item()

        # 记录训练损失
        train_losses.append(loss_temp / len(train_dataloader))

        # 测试阶段
        My_VGG16.eval()  # 设置为评估模式
        total_test_loss = 0
        total_accuracy = 0
        with torch.no_grad():
            for data in val_dataloader:
                imgs, targets = data
                imgs = imgs.to(device)
                targets = targets.to(device)
                outputs = My_VGG16(imgs)
                loss = loss_fn(outputs, targets)

                total_test_loss += loss.item()
                #找到每个样本的预测类别，然后与真实标签进行比较
                accuracy = (outputs.argmax(1) == targets).sum().item()  # 使用.item()将Tensor转换为Python数值
                total_accuracy += accuracy   #计算正确预测的数量并累加到 total_accuracy

imgs, targets = data这段代码

可以从图中看出data有两个部分分别是(8,3,224,224),(8,) 一定要注意顺序分别对应的就是imgs，targets

outputs = My_VGG16(imgs)  outputs = tensor（8，3）

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

  # 记录验证损失和准确率
        val_losses.append(total_test_loss / len(val_dataloader))
        accuracies.append(total_accuracy / val_data_size)

        print("整体测试集上的正确率:{}".format(total_accuracy / val_data_size))
        print("整体测试集上的Loss：{}".format(total_test_loss))

    # 保存模型
    torch.save(My_VGG16, "vgg16_{}.pth".format(2))
    torch.save(My_VGG16.state_dict(), "vgg16_dict_{}.txt".format(2))
    print("模型已保存")

    # 绘制损失和准确率曲线
    epochs_range = range(1, epochs + 1)

    plt.figure(figsize=(12, 5))  #创建一个新的图形窗口，设置图形的大小

    # 训练和验证损失
    plt.subplot(1, 2, 1)  #第一个图像框
    plt.plot(epochs_range, train_losses, label='Training Loss')
    #train_losses=[1.0952595781396937, 1.085733965591148, 1.07982603708903, 1.0776626953372248, 1.0816458507820412]
    plt.plot(epochs_range, val_losses, label='Validation Loss')
    plt.title('Loss vs Epochs')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()

    # 准确率
    plt.subplot(1, 2, 2)
    plt.plot(epochs_range, accuracies, label='Validation Accuracy')
    plt.title('Accuracy vs Epochs')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.legend()

    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    parse_opt()
    train()
    # 测试模型
    # model = VGG16()
    # sample_input = torch.randn(1, 3, 224, 224)  # Batch size 为 1
    # output = model(sample_input)
    # print("Final output shape:", output.shape)  # 应输出 [1, 3]

代码讲解，注释很详细，不做细数。

输出图片效果，这是只训练了5次的网络，可以看到训练效果是不理想的。

以上就是训练集代码及讲解。
训练整体代码：

#VGG_16训练网络模型，第一代测试代码，主要功能如下
#识别3类蚂蚁，蜜蜂，狗，可任意添加，文件名为1.jpg,1.txt,内容ants
#2代是在一代的基础上添加了图像显示，标准化参数设计方面
import os
import parser
import time
import numpy as np
import torch
import torchvision.transforms as transforms
import PIL
from PIL import Image
from scipy.stats import norm
from torch import nn
from torch.utils.data import Dataset, random_split, DataLoader, ConcatDataset
import torchvision.datasets
from torch import nn
from torch.nn import MaxPool2d, Flatten

import argparse
from pathlib import Path


# ------------------------------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 opt

opt = parse_opt()  # 调用解析函数
epochs = opt.epochs              # 训练的轮数
batch_size = opt.batch_size      # 每个批次的样本数量
learning_rate = opt.learning_rate
device = torch.device(opt.device)

class CustomDataset(Dataset):
    def __init__(self, img_dir, label_dir, transform=None):
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.transform = transform
        self.img_labels = self.load_labels()

    def load_labels(self):
        label_map = {'ants': 0, 'bees': 1,'dogs': 2}    #字典label_map 来定义标签与数字之间的映射关系
        labels = []                                     #创建一个空列表 labels，用来存储从文件中读取到的标签的数字编码
        for label_file in os.listdir(self.label_dir):
            with open(os.path.join(self.label_dir, label_file), 'r') as f:  #对于每个标签文件，使用 open() 打开文件并读取第一行内容readline()
                line = f.readline().strip()         #strip() 方法用于去掉行首和行尾的空白字符（包括换行符）
                #从label_map 中获取当前行的标签对应的数字编码。如果当前行的内容不在 label_map 中，则返回-1
                labels.append(label_map.get(line, -1))
        return labels

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        #使用 os.path.join 将图像目录路径 (self.img_dir) 和图像文件名拼接起来。文件名格式为 idx + 1（即从1开始计数），加上 .jpg 后缀。
        #例如，如果idx为0，则img_path会是 D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_image\\1.jpg
        img_path = os.path.join(self.img_dir, str(idx+1) + '.jpg')
        image = Image.open(img_path).convert('RGB')  #打开指定路径的图像文件，并将其转换为 RGB 模式。
        label = self.img_labels[idx]   #列表中获取当前图像的标签，idx 是当前图像的索引

        if self.transform:   #transform变量是在CustomDataset中最终调用的transforms.Compose实现
            image = self.transform(image)

        return image, label

# 定义转换，包括调整大小
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整为 224x224 大小
    transforms.ToTensor(),          # 转换为 Tensor
    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    # 可以添加其他转换，例如归一化等
])

# 初始化数据集
ants_data = CustomDataset(
    img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_image',
    label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\ants_label',
    transform=transform
)


bees_data = CustomDataset(
    img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\bees_image',
    label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\bees_label',
    transform=transform
)

dogs_data = CustomDataset(
    img_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\dogs_image',
    label_dir='D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\train\\dogs_label',
    transform=transform
)
#-----------------------2.将数据集总和并随机打乱并分成训练集测试集-------------------
# 合并数据集
total_data = ants_data+bees_data
total_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_size
val_data_size = val_size

print("训练集长度：{}".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):
    def __init__(self):
        super(VGG16, self).__init__()
        self.model = nn.Sequential(  # 构建子类，可以减少foward代码撰写
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), #最大池化将224变成112
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), # 56x56
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 28x28
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), # 14x14
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 7x7
        )
        # 使用AdaptiveAvgPool2d来确保输出为固定大小
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))

        # 全连接层
        self.fc1 = nn.Linear(in_features=512 * 7 * 7, out_features=4096)
        self.fc2 = nn.Linear(in_features=4096, out_features=4096)
        self.fc3 = nn.Linear(in_features=4096, out_features=3)

        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.3)

    def forward(self, x):
        x = self.model(x)
        x = self.avgpool(x)  # 使用自适应平均池化
        x = x.view(x.size(0), -1)  # 拉平操作,(batch_size, num_features) 也就是（8，25088） 512*7*7 = 25088
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x

#------------------------------------4.训练加测试-----------------------

import numpy as np
import matplotlib
matplotlib.use('TkAgg')  # 或者尝试 'Qt5Agg'，有这行代码会多一个弹窗显示
import matplotlib.pyplot as plt
import torch.nn as nn


def train():
    My_VGG16 = VGG16()              #调用类
    My_VGG16 = My_VGG16.to(device)

    # (1).损失函数构建
    loss_fn = nn.CrossEntropyLoss()  #计算预测值与真实标签之间的差异
    loss_fn = loss_fn.to(device)     #将模型和数据都放在同一个设备上，GPU

    # (2).优化器
    # #随机梯度下降（Stochastic Gradient Descent）优化器的一种实现。SGD 是一种常见的优化算法
    optimizer = torch.optim.SGD(My_VGG16.parameters(), lr=learning_rate)


    # 用于存储损失和准确率
    train_losses = []
    val_losses = []
    accuracies = []

    for i in range(epochs):
        loss_temp = 0  # 临时变量
        print("--------第{}轮训练开始--------".format(i + 1))

        # 训练阶段
        My_VGG16.train()  # 设置为训练模式
        #data有两个部分分别是(8,3,224,224),(8,) 一定要注意顺序
        for data in train_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)
            #outputs = tensor（8，3）
            outputs = My_VGG16(imgs)
            loss = loss_fn(outputs, targets)

            # 优化器优化模型
            optimizer.zero_grad()  # 梯度清零
            loss.backward()  # 反向传播
            optimizer.step()  # 梯度更新

            loss_temp += loss.item()

        # 记录训练损失
        train_losses.append(loss_temp / len(train_dataloader))

        # 测试阶段
        My_VGG16.eval()  # 设置为评估模式
        total_test_loss = 0
        total_accuracy = 0
        with torch.no_grad():
            for data in val_dataloader:
                imgs, targets = data
                imgs = imgs.to(device)
                targets = targets.to(device)
                outputs = My_VGG16(imgs)
                loss = loss_fn(outputs, targets)

                total_test_loss += loss.item()
                #找到每个样本的预测类别，然后与真实标签进行比较
                accuracy = (outputs.argmax(1) == targets).sum().item()  # 使用.item()将Tensor转换为Python数值
                total_accuracy += accuracy   #计算正确预测的数量并累加到 total_accuracy
#-----------------------5.保存数据至.pth文件并绘图---------------------------------------------------
        # 记录验证损失和准确率
        val_losses.append(total_test_loss / len(val_dataloader))
        accuracies.append(total_accuracy / val_data_size)

        print("整体测试集上的正确率:{}".format(total_accuracy / val_data_size))
        print("整体测试集上的Loss：{}".format(total_test_loss))

    # 保存模型
    torch.save(My_VGG16, "vgg16_{}.pth".format(2))
    torch.save(My_VGG16.state_dict(), "vgg16_dict_{}.txt".format(2))
    print("模型已保存")

    # 绘制损失和准确率曲线
    epochs_range = range(1, epochs + 1)

    plt.figure(figsize=(12, 5))  #创建一个新的图形窗口，设置图形的大小

    # 训练和验证损失
    plt.subplot(1, 2, 1)  #第一个图像框
    plt.plot(epochs_range, train_losses, label='Training Loss')
    #train_losses=[1.0952595781396937, 1.085733965591148, 1.07982603708903, 1.0776626953372248, 1.0816458507820412]
    plt.plot(epochs_range, val_losses, label='Validation Loss')
    plt.title('Loss vs Epochs')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()

    # 准确率
    plt.subplot(1, 2, 2)
    plt.plot(epochs_range, accuracies, label='Validation Accuracy')
    plt.title('Accuracy vs Epochs')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.legend()

    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    parse_opt()
    train()
    # 测试模型
    # model = VGG16()
    # sample_input = torch.randn(1, 3, 224, 224)  # Batch size 为 1
    # output = model(sample_input)
    # print("Final output shape:", output.shape)  # 应输出 [1, 3]

注意本人的训练是用的gpu，不推荐用cpu跑代码训练。

检测部分代码撰写

这一部分的代码主要实现的是输入任意大小的图片，预测出它的种类。

1.导入头文件参数初始化
导入头文件，参数标准化后，读取图片，并通过transform实现图片输入格式的设置。

import torch
import torchvision
from PIL import Image
from torch import nn
from torch.nn import MaxPool2d
import argparse
from pathlib import Path

#--------------------------1.参数数据初始化-----------------------------------
def parse_opt():
    parser = argparse.ArgumentParser()  # 创建 ArgumentParser 对象
    parser.add_argument('--Dropout_Val', type=int, default=0.4, help='Dropout_Val number')  # 添加参数
    #--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 opt


opt = parse_opt()  # 调用解析函数
Dropout_Val = opt.Dropout_Val
device = torch.device(opt.device)


image_path = "D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\1.jpg"
image = Image.open(image_path)   #注意格式为RGBA
print(image)

image = image.convert("RGB")     #将格式变成RGB
print(image)

#确保你的图像预处理步骤（如归一化）与模型训练时使用的步骤相匹配。通常，VGG16 需要将图像归一化到 [0, 1] 范围内，
# 并可能需要减去均值和除以标准差
transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((224, 224)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

image = transform(image)
print(image.shape)

# 类别标签
class_names = ['ant', 'bee', 'dog']  # 根据实际情况填写

2.搭建VGG16网络结构

class VGG16(nn.Module):
    def __init__(self):
        super(VGG16, self).__init__()
        self.model = nn.Sequential(  # 构建子类，可以减少foward代码撰写
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), #最大池化将224变成112
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), # 56x56
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 28x28
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), # 14x14
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 7x7
        )
        # 使用AdaptiveAvgPool2d来确保输出为固定大小
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))

        # 全连接层
        self.fc1 = nn.Linear(in_features=512 * 7 * 7, out_features=4096)
        self.fc2 = nn.Linear(in_features=4096, out_features=4096)
        self.fc3 = nn.Linear(in_features=4096, out_features=3)

        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(Dropout_Val)

    def forward(self, x):
        x = self.model(x)
        x = self.avgpool(x)  # 使用自适应平均池化
        x = x.view(x.size(0), -1)  # 拉平操作
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x

3.输出测试结果
输出测试结果需要将图片送入VGG16网络进行计算，带入训练好的权值vgg16_1.pth，得出预测结果，将预测的结果输出成概率从而计算是什么品种。

model = VGG16()
model = torch.load("vgg16_1.pth")
#模型的权重在GPU上，而输入张量在CPU上。要解决这个问题，你可以将输入张量移动到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
image = image.to(device)
#模型的权重在CPU上，而输入张量在GPU上。要解决这个问题，你可以将输入张量移动到CPU
#model = torch.load("tudui_0.pth",map_location=torch.device('cpu'))
image = torch.reshape(image,(1,3,224,224))
model.eval()
with torch.no_grad():
    output = model(image)
predicted_index = output.argmax(1).item()  # 获取预测的索引
print(output.argmax(1))
print(output)   #输出结果预测的是第六个类别
print("Predicted index:", predicted_index)
# 输出对应的类别名称
print("Predicted class:", class_names[predicted_index])

测试输出效果

其实多训练几次预测效果还是可以的。

测试整体代码

# vgg16网络测试文件,第一代测试代码，主要功能如下
# 任意输入一张图片，输出判断类型,二代与一代相比就是多了参数标准化

import torch
import torchvision
from PIL import Image
from torch import nn
from torch.nn import MaxPool2d
import argparse
from pathlib import Path

#--------------------------1.参数数据初始化-----------------------------------
def parse_opt():
    parser = argparse.ArgumentParser()  # 创建 ArgumentParser 对象
    parser.add_argument('--Dropout_Val', type=int, default=0.4, help='Dropout_Val number')  # 添加参数
    #--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 opt


opt = parse_opt()  # 调用解析函数
Dropout_Val = opt.Dropout_Val
device = torch.device(opt.device)


image_path = "D:\\Pycharm\\Pytorch_test\\数据集\\练手数据集\\1.jpg"
image = Image.open(image_path)   #注意格式为RGBA
print(image)

image = image.convert("RGB")     #将格式变成RGB
print(image)

#确保你的图像预处理步骤（如归一化）与模型训练时使用的步骤相匹配。通常，VGG16 需要将图像归一化到 [0, 1] 范围内，
# 并可能需要减去均值和除以标准差
transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((224, 224)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

image = transform(image)
print(image.shape)

# 类别标签
class_names = ['ant', 'bee', 'dog']  # 根据实际情况填写

#-------------------------------------2.搭建VGG16网络结构---------------------------------------------
class VGG16(nn.Module):
    def __init__(self):
        super(VGG16, self).__init__()
        self.model = nn.Sequential(  # 构建子类，可以减少foward代码撰写
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), #最大池化将224变成112
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), # 56x56
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 28x28
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2), # 14x14
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 7x7
        )
        # 使用AdaptiveAvgPool2d来确保输出为固定大小
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))

        # 全连接层
        self.fc1 = nn.Linear(in_features=512 * 7 * 7, out_features=4096)
        self.fc2 = nn.Linear(in_features=4096, out_features=4096)
        self.fc3 = nn.Linear(in_features=4096, out_features=3)

        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(Dropout_Val)

    def forward(self, x):
        x = self.model(x)
        x = self.avgpool(x)  # 使用自适应平均池化
        x = x.view(x.size(0), -1)  # 拉平操作
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x

#---------------------------3.输出测试结果-----------------------------
model = VGG16()
model = torch.load("vgg16_1.pth")
#模型的权重在GPU上，而输入张量在CPU上。要解决这个问题，你可以将输入张量移动到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
image = image.to(device)
#模型的权重在CPU上，而输入张量在GPU上。要解决这个问题，你可以将输入张量移动到CPU
#model = torch.load("tudui_0.pth",map_location=torch.device('cpu'))
image = torch.reshape(image,(1,3,224,224))
model.eval()
with torch.no_grad():
    output = model(image)
predicted_index = output.argmax(1).item()  # 获取预测的索引
print(output.argmax(1))
print(output)   #输出结果预测的是第六个类别
print("Predicted index:", predicted_index)
# 输出对应的类别名称
print("Predicted class:", class_names[predicted_index])

最开始训练的次数有点少，导致预测结果不理想，后来去吃饭前设置训练epochs = 500次，回来发现结果变好了一点。提高准确率的方式很多，还可以丰富训练集。注意我的权值文件权值vgg16_1.pth，可能使用的是权值vgg16_2.pth，代码我就不改了。主要是要理解以上代码。

以上就是本人的心得与总结，如有不足之处请多多包涵。

百度网盘链接代码权值及训练图片： https://pan.baidu.com/s/14xixodOGVEMK-KpE5rZTtw

提取码: dwrg