YOLOv4-tiny（一）模型构建与训练实现

itoN1ght

已于 2024-05-10 11:20:00 修改

阅读量1.7k

点赞数 47

文章标签： YOLO python 目标检测图像处理神经网络深度学习 pytorch

于 2024-05-06 09:30:33 首次发布

本文链接：https://blog.csdn.net/Wexssss/article/details/138482805

版权

YOLOv4-tiny（一）模型构建与训练实现

Tags: YOLO, 深度学习, 目标检测
Created: April 24, 2024 3:20 PM
Last edited time: May 6, 2024 9:15 AM
Status: Done

模型架构图

在这里插入图片描述

零. 引言

YOLOv4-tiny模型是在YOLOv4模型基础上优化裁剪的轻量化模型，模型参数量只有600万（YOLOv4有6000万），因此该模型速度非常快，准确率也还可以，对于想入门学习的小白（包括我自己，哈哈）我个人觉得还是比较合适的。我个人的学习路线是：1.下载模型源码 2.部署运行环境 3.进行推理测试 4.关键步骤：对照模型架构使劲撸源码（PS.中途恶补了两本书，外行亚历山大）。我是纯外行，所以感觉先调通模型会比较有成就感，虽然基本上一行代码都看不懂，但是当视频中我的脑袋被框框检测到的时候还是很开心的。借着这篇文章算是一个学习记录，也希望能帮助到有需要的人。（~~虽然是外行，但是我是认认真真的在学的，有错误的地方欢迎批评指正。~~）

再碎碎念几句！

我个人的体会是，目标检测实际上是在图片上进行目标的定位和分类。图片的本质就是像素点，像我们常见的1920x1080尺寸的图片，实际上就是宽高方向各这么多数量的像素点，每张图片一般是RGB（红绿蓝）三个颜色通道，每个像素点上是0~255的数值，代表这个像素在本通道的颜色取值。因此，在对图片进行目标检测时，实际上输入模型的就是这些数值。因此，模型本身还是对数据进行操作变换。其实对我来说，直接看这种成熟的模型还是有压力的，所以我先去了解了下更入门的MNIST数据的检测模型，这个检测模型就是对**‘手写数字0-9’**的检测，图片也是白底黑字，非常简单，模型的目的是将手写数字图片输入，然后判断写的是哪个数字，本质上就是个10分类问题。有兴趣的可以先看看，可以直观了解下神经网络模型的运行机理和常见概念。

为了理解起来更纯粹一些，我将源码中的关于训练train.py使用到代码单独摘出来，其他的可视化、评估、推理等模块此文暂时不讲，只针对训练部分所需要的内容进行解释。

另外，我个人的笨办法是，撸源码的时候建议先看，然后一定一定要自己一行一行编一下，熟悉模型算法的同时，还能学习下python。

出于某种神秘的指引，我感觉小白的学习经验更能适合小白，你笑话我我也不怕，毕竟大家都挺菜的。 /手动狗头。

感谢Bubbliiiing导的源码，我的工作主要是在原有的注释基础上，做了大量入门级的注释，并将训练部分单独摘出来，方便理解，新手可以先看本文的源码。

先贴上源码，如果有不懂的可以评论区留言。

本文的源码：https://github.com/iton1ght/Yolov4-tiny-clone/tree/main

整个项目的源码：https://github.com/bubbliiiing/yolov4-tiny-pytorch

一. 模型架构图解释

1.1 模型结构组成（程序组成）

模型结构整体放到nets文件夹中，包括三个程序，net_yolo.py、CSPdarknet53_tiny.py、net_loss.py。其中前两个程序存放模型的模块或层，最后一个程序存放损失模块以及初始化和学习率的函数

net_yolo.py：包括三个模块，一个函数
- BasicConv：定义了卷积标准化激活模块
- Upsample：定义了上采样模块
- yolo_head：定义头部函数
- YoloBody：定义了整个模型模块
CSPdarnnet53_tiny.py：包括三个模块，一个函数
- BasicConv：定义了卷积标准化激活模块
- Resblock_body：定义了残差模块
- CSPDarkNet：定义了CSPDarkNet模块
- darknet53_tiny：定义函数，函数功能在CSPDarkNet模块基础上定义了一个实例，并判断是否加载主干的预训练权重
net_loss.py：定义了一个模块，三个函数
- YoloLoss：定义损失模块
- weights_init：定义权重初始化函数
- get_lr_scheduler：定义函数用于获取学习率调度器
- set_optimizer_lr：定义函数用于设置优化器的学习率

在这里插入图片描述

关于卷积部分的基础知识，可以看下面这篇文章，先了解卷积在图片上是如何操作的，以及卷积的基本参数概念。

深度学习基础入门篇[9.1]：卷积之标准卷积：卷积核/特征图/卷积计算、填充、感受视野、多通道输入输出、卷积优势和应用案例讲解

1.2 模型计算流程

1.2.1 模型前向传播流程

简单解释下前向传播，直观理解就是将图片输入，按照模型的流程对图片进行处理（处理像素值），最终输出两组特征图。处理方式有许多种，在该模型中包括卷积、批量归一化、激活函数、池化、上采样等，目的是提取图片的特征。

模型输入：416x416x3 → 代表输入416x416像素的图片，图片为3通道数RGB
backbone：模型的主干特征提出网络，用于提取图片的特征。
- backbone主体由三个resblock_body模块组成，另外还包括几个卷积层
- 输入的图片经处理后，输出为两路，一路为经backbone全部计算输出，另一路由第三个resblock_body模块中间引出。
yolo_head：模型的头部，本模型由两个，输入是由经backbone输出的两路再经一定处理后，输出为模型的输出
模型输出：13x13xc和26x26xc → 输出为两组特征图，c代表通道数，与锚框数量、分类类别有关。

1.2.2 模型计算损失流程

简单解释下损失，所谓的损失本质上是误差，即目标值和预测值之间的差距。这个差距是在特征图尺度上进行计算的，也就是说对于图片来说，a（不带标注）和b（带标注），a和b都要转化到特征图尺度上，a为预测值，b为真实值。两者在特征图尺度上的误差即为损失。这个损失一方面可以评估模型的准确性，另一方面是为了反向传播使用的，最终的损失要转化到各模块或层的损失，从而得到权重的梯度下降值，完成权重的更新。

针对输出得到的两组特征图数据，并结合先验框得到相应的预测框，与原始图的真实框进行计算，得到预测框和真实框相应的损失loss

1.2.3 模型反向传播流程

反向传播的方法在pytorch框架已经定义，直接调用即可，无需纠结。

直接调用Pytorch中的反向传播方法
利用优化器更新权重

二. 模型主体（包括各类模块和卷积层）

2.1 卷积标准化激活模块（ConvBNLeaky=Conv+BN+Leaky）

卷积标准化激活模块由1个标准二维卷积层、1个批量归一化层、1个激活函数层组成。

另外给小伙伴解释下图中的参数含义：

k3是kernel_size=3缩写，代表卷积核为3x3尺寸
s2是stride=2缩写，代表卷积核在图片上每次移动2个像素点
c32是out_channels=32的缩写，代表经卷积核处理后，输出通道数为32个。（为什么不是输入通道呢，因为输入通道数是确定的数量，等于上一个模块的输出通道数。例如对图片输入的第一个卷积模块来说，由于图片的通道是3，该卷积模块的输入通道就是3）

另外还得给新手小伙伴解释下nn.Module父类的特点，通过继承该父类，可以定义forward方法，就是说可以通过实例直接调用该方法，从而实现前向传播。

在这里插入图片描述

# 定义基础卷积块:Conv2d + BatchNorm2d + LeakyReLU
# 在PyTorch中，当你定义一个nn.Module的子类时，你通常实现一个forward方法，该方法定义了模型的前向传播。
# 当你创建这个类的实例并调用它时，PyTorch在背后实际上调用了forward方法。这就是为什么你可以使用类似model(input)的
# 语法来执行模型的前向传播，即使model的类定义中并没有__call__方法。
class BasicConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, kernel_size//2, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)
        self.activation = nn.LeakyReLU(0.1)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.activation(x)
        return x

2.2 上采样模块（ConvBNLeaky+Upsample）

上采样模块由一个卷积标准化激活模块和一个上采样层组成，通过Sequential序列容器实现顺序执行。

上采样顾名思义，就是通过上采样操作，改变图片的尺寸，实际上就是通过插值或者邻近值等方式增加像素点。

在这里插入图片描述

# 定义上采样模块
# 上采样模块包括：ConvBNLeady + Unsample
# nn.Sequential 是 PyTorch 中用于创建一个包含多个模块的序列容器。这个容器会按照模块在构造函数中传入的顺序来执行它们
class Upsample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.upsample = nn.Sequential(
            BasicConv(in_channels, out_channels, kernel_size=1),
            nn.Upsample(scale_factor=2, mode='nearest')
        )
    def forward(self, x,):
        x = self.upsample(x)
        return x

2.3 Resblock_body模块（ConvBNLeaky+Maxpool）

Resblock_body模块由四个卷积标准化激活模块和一个二维最大池化层组成，中间引出一大一小残差边，其中第一个大残差边为第一个卷积标准化激活模块的输出（通道数不变），此处还运用torch.split方法对通道进行分割，只有一半通道进入下面的运算。第二个残差边为第二个卷积标准化激活模块的输出。模块里还运用了torch.cat方法，进行通道的合并。

在这里插入图片描述

# 定义Resblock_body模块
# 模块的特点为存在一大一小残差边，大残差边为第一次卷积标准化激活的输出，小残差边为第二次卷积标准化激活的输出
class Resblock_body(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        # 定义该模块的输入通道数和输出通道数属性，对一个实例来说，输入为64，输出为128
        self.in_channels  = in_channels
        self.out_channels = out_channels
        # 卷积层
        self.conv1 = BasicConv(self.in_channels, self.in_channels, 3)
        self.conv2 = BasicConv(self.in_channels//2, self.in_channels//2, 3)
        self.conv3 = BasicConv(self.in_channels//2, self.in_channels//2, 3)
        self.conv4 = BasicConv(self.out_channels//2, self.out_channels//2, 1)

        # 二维最大池化层
        self.maxpool = nn.MaxPool2d([2,2], [2,2])
    def forward(self, x):
        # 第一个卷积3x3
        x = self.conv1(x)
        # 引出大残差边
        route_1 = x
        # 对特征层的通道进行分割，取第二部分进行主干运算，第二部分通道数为原来的一半
        c = self.in_channels
        x = torch.split(x, c//2, 1)[1] # 沿着第二个维度（通道维度）将张量 x 分割成大小为 c//2 的块，并取这些块中的第二个块作为新的 x。
        # 第二个卷积3x3
        x = self.conv2(x)
        # 引出小残差边
        route_2 = x
        # 第三个卷积3x3
        x = self.conv3(x)
        # 小残差边与主干进行通道维度合并
        x = torch.cat([x, route_2],1)
        # 第四个卷积1x1，卷积核为1时，通常为改变通道数
        x = self.conv4(x)
        # 输出特征层
        feat = x
        # 大残差边与主干进行通道维度合并
        x = torch.cat([x,route_1],1)
        # 进行池化层操作,对高和宽进行压缩
        x = self.maxpool(x)
        return x, feat

2.4 CSPdarknet模块（ConvBNLeaky+Resblock_body）

CSPdarknet模块实际上就是该模型的主干网络结构，也就是是BackBone，主要是负责图片的特征提取，由前面建立的ConvBNLeaky模块和Resblock_body模块构建，包括三个ConvBNLeaky模块和三个Resblock_body模块。BackBone模块输出两个特征图，分别为feat1和feat2，经过一定处理后再输入到头模块中，完成结果的输出。

模块里进行了主干网络结构的权重初始化。这是为了在后续darknet53_tiny函数中进行权重加载。

#定义CSPdarknet模块
class CSPDarkNet(nn.Module):
    def __init__(self):
        super().__init__()

        # 两个标准化卷积模块
        # 第一个卷积3x3，416x416x3 -> 208x208x32
        self.conv1 = BasicConv(3, 32, 3,2)
        # 第二个卷积3x3, 208x208x32 -> 104x104x64
        self.conv2 = BasicConv(32, 64,3 ,2)
        # 三个resblock_body模块
        # 104x104x64 -> 52x52x128
        self.resblock_body1 = Resblock_body(64, 128)
        # 52x52x128 -> 26x26x256
        self.resblock_body2 = Resblock_body(128, 256)
        # 26x26x256 -> 13x13x512
        self.resblock_body3 = Resblock_body(256, 512)
        # 一个标准化卷积模块
        # 第三个卷积3x3，13x13x512 -> 13x13x512
        self.conv3 = BasicConv(512,512,3,1)

        # 定义初始化权重
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 
                m.weight.data.normal_(0, math.sqrt(2./n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def forward(self, x):

        x = self.conv1(x)
        x = self.conv2(x)
        x, _ = self.resblock_body1(x)
        x, _ = self.resblock_body2(x)
        x, feat1 = self.resblock_body3(x) #feat1 为26x26x256
        x = self.conv3(x)
        feat2 = x  #feat2 为13x13x512
        return feat1, feat2

# 主干网络结构已经进行了权重初始化，这个函数建立了一个CSPDarkNet实例，并判断是否加载预训练权重，并返回该实例
def darknet53_tiny(pretrained,**kwargs):
    model = CSPDarkNet()
    # 如果pretrained为TRUE,则加载主干的预训练权重。若已经全模型加载了预训练权重，则pretrained的值无所谓，主干部分权重已经在全模型预训练权重中被加载了
    if pretrained:
        model.load_state_dict(torch.load("model_data/CSPdarknet53_tiny_backbone_weights.pth"))
    return model

2.5 yolo_head头部方法（函数）

模型结构中含两路输出，因此有两个头部方案，方法的逻辑是一样的，都是经卷积标准化激活模块后，在经二维标准卷积进行输出。最终输出的通道数c与分类数量有关，因此这里用变量c代表

c=len(anchors_mask[l])*(5+classes_num)，其中len()代表锚框数量，5代表锚框的xywhc五个参数，classes_num代表类别数量

在这里插入图片描述

# 定义头部模块
# filters_list是一个包括两个元素的列表，分别代表最后一个标准卷积的输入通道数和输出通道数
# 第一个元素代表输入通道数（512或256），第二个元素代表c，也就是输出通道数
def yolo_head(filters_list, in_filters):
    m = nn.Sequential(
        BasicConv(in_filters, filters_list[0], kernel_size=3),
        nn.Conv2d(filters_list[0], filters_list[1], kernel_size=1, stride=1)
    )
    return m

2.6 Yolo_Body

前面所有的模块和层搭建好后，开始进行主体的组建。这部分内容可以对照模型的整个结构图来看。

class YoloBody(nn.Module):
    def __init__(self, anchor_mask, num_classes, phi=0, pretrained=False):
        super().__init__()

        self.phi = phi
        self.backbone = darknet53_tiny(pretrained)

        self.conv_for_P5  = BasicConv(512, 256, 1, 1)
        self.yolo_headP5 = yolo_head([512, len(anchor_mask[0])*(5+num_classes)], 256)

        self.upsample     = Upsample(256, 128)
        self.yolo_headP4 = yolo_head([256, len(anchor_mask[1])*(5+num_classes)], 384)

    def forward(self, x):

        feat1, feat2 = self.backbone(x)

        P5   = self.conv_for_P5(feat2)
        out0 = self.yolo_headP5(P5)

        P5_Upsample = self.upsample(P5)
        P4   = torch.cat([P5_Upsample, feat1], 1)
        out1 = self.yolo_headP4(P4)

        return out0, out1

2.7 模型搭建的一些注意事项

各种模块或层的名称不要改动，所有的权重都是以字典形式存储，其中模块名或层名为键、权重值为值，在加载训练好的权重时，是通过‘键’来匹配并加载的，如果改动了，则会匹配不上。
搭建模型的过程先不要考虑为什么这么搭建的问题，优先考虑如何代码实现的问题，模型都是前人测试过的，虽然不是最新的（目前最新是YOLOv8），但是工程化应用还是可以的。

三. 损失模块

损失模块是比较难理解的部分，涉及先验框、预测框和真实框等算法概念，以及交并比、MSE损失、BSE损失等评估方式。

损失模块定义了一个class YoloLoss(nn.Module)的类，在该类中定义了模块的基本属性和各种方法，包括以下：

3.1 辅助方法（简单）

def __init__(self, anchors, num_classes, input_shape, cuda, anchors_mask=[[6, 7, 8], [3, 4, 5], [0, 1, 2]], label_smoothing=0)：定义了类初始化方法，包括各种参数初始化
def clip_by_tensor(self, t, t_min, t_max)：定义了张量裁剪方法，将张量的每个元素裁剪至[t_min,t_max]区间
def MSELoss(self, pred, target)：定义MSELoss损失函数，预测值与目标值的差值的平方
def BSELoss(self, pred, target)：定义BSELoss损失函数，二元交叉熵，注意二院交叉熵的常用于二分类问题，也就是说target中的元素要么是0，要么是1
def smooth_labels(self, y_ture, label_smoothing, num_classes)：该方法用于平滑真实的标签，略小于1略大于0

3.2 核心方法（困难）

3.2.1 def box_ciou(self, box_1, box_2)

该方法是计算真实框张量和预测框张量的所有ciou，输入张量的形状均为五维，包括批次数、先验框数量、特征图宽、高、以及真实框/预测框的xywh，返回值为所有的ciou。这里的输入为张量，可以理解为每个先验框下的每张图片的每个像素点均有预测框和真实框，数量为batch x anchor_num x feat_w x feat_h

方法先根据框的xywh求的框的左上角和右上角坐标，然后由此求的真实框和先验框的交集和并集，再求得交并比iou
将两个中心距离引入交并比，中心距离越小，ciou越大，这里利用最小包络框的对角线距离进行归一化操作（对角线效果比固定值要好，因为两者是同向变化，可以减缓中心距离对ciou的影响）
将真实框和预测框的宽长比引入交并比，从而体现两个框的形状重合特性。先计算真实框和预测框的宽长比反正切值的差值的平方，在设置损失权重alpha，用于调节宽长比在ciou的比例，若交并比iou较小，则权重alpha较大，宽长比对ciou的影响增大，会使ciou快速降低，原因是若无重合则形状相似也无用

方法源码：

def box_ciou(self, box_1, box_2):
        """
        函数定义：计算真实框张量和预测框张量的所有CIoU
        输入为：
        ----------
        box_1: tensor, shape=(batch, anchor_num, feat_w, feat_h, 4), xywh
        box_2: tensor, shape=(batch, anchor_num, feat_w, feat_h, 4), xywh

        返回为：
        -------
        ciou: tensor, shape=(batch, anchor_num, feat_w, feat_h, 1)
        """
        # 求真实框左上角和右下角坐标
        box_1_xy   = box_1[..., :2]
        box_1_wh   = box_1[..., 2:4]
        box_1_mins = box_1_xy - box_1_wh / 2.
        box_1_maxs = box_1_xy + box_1_wh / 2.

        # 求预测框左上角和右下角坐标
        box_2_xy   = box_2[..., :2]
        box_2_wh   = box_2[..., 2:4]
        box_2_mins = box_2_xy - box_2_wh / 2.
        box_2_maxs = box_2_xy + box_1_wh / 2.

        # 求真实框和预测框‘所有的’交集的左上角和右下角坐标，并求交集面积和并集面积，求得交并比
        intersect_mins = torch.max(box_1_mins, box_2_mins)
        intersect_maxs = torch.min(box_1_maxs, box_2_maxs)
        intersect_wh   = torch.max(intersect_maxs - intersect_mins, torch.zeros_like(intersect_maxs))
        intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
        box_1_area     = box_1_wh[..., 0] * box_1_wh[..., 1]
        box_2_area     = box_2_wh[..., 0] + box_2_wh[..., 1]
        union_area     = box_1_area + box_2_area - intersect_area
        iou            = intersect_area / torch.clamp(union_area, 1e-6)

        # 求两个框中心点距离，为了简化运算避免开方，则求取中心点距离的平方
        center_distance = torch.sum(torch.pow(box_1_xy - box_2_xy, 2), dim=-1)  # 此处指定最后一维为计算对象，center-distance的张量形状与iou一致

        # 计算包裹真实框和预测框的最小包络框,求取包络框的左上角和右下角
        enclose_mins   = torch.min(box_1_mins, box_2_mins)
        enclose_maxs   = torch.max(box_1_maxs, box_2_maxs)
        enclose_wh     = torch.max(enclose_maxs - enclose_mins, torch.zeros_like(enclose_maxs))

        # 求包络框的对角线距离
        enclose_diagonal = torch.sum(torch.pow(enclose_wh, 2), dim=-1)

        # 将中心线比例引入交并比iou，center_distance越小,ciou越大
        # 其中利用包络框的对角线距离对center_distance进行归一化操作
        # （注：也可以用其中固定值进行归一化，感觉用enclose_diagonal进行归一化，可以‘减缓’center_distance参数对ciou的影响，因为center_distance和enclose_diagonal是同向变化的）
        ciou = iou - 1.0 * center_distance / torch.clamp(enclose_diagonal, 1e-6)

        # 将真实框和预测框的宽长比引入交并比，从而体现两个框的形状重合特性
        box_wh_ratio = (4 / (math.pi ** 2)) * torch.pow((torch.atan(box_1_wh[..., 0] / torch.clamp(box_1_wh[..., 1], 1e-6)) - torch.atan(box_2_wh[..., 0] / torch.clamp(box_2_wh[..., 1], 1e-6))), 2)
        # 系数alpha用于调整宽高比损失项的权重，当iou减小时，alpha增大，代表增大权重，使得ciou快速降低，因为若无重合，即使形状相似也无用
        # alpha是关于box_wh_ratio的增函数，此处为归一化操作
        alpha = box_wh_ratio / torch.clamp((1.0 - iou + box_wh_ratio), 1e-6)
        ciou = ciou - alpha * box_wh_ratio
        return ciou

3.2.2 def box_iou(self, box_a, box_b)

该方法是计算一张图片上的所有真实框和先验框的iou，输入张量的形状均为二维，包括真实框/先验框数量，以及相应的框的xywh，返回值为所有的iou

该方法与上述方法的计算方式基本相同
该方法主要对一张图片来说，计算图片上所有真实框和先验框的交并比，从而选出对每一个真实框来说，最匹配的先验框，并获取其序号
该方法后续将用于方法2.3和方法2.4，在2.3中用于获得真实目标张量的最匹配先验框序号，真实目标张量为y_true = torch.zeros(bs, anchors_num, in_h, in_w, self.bbox_attrs, requires_grad=False)，第二维是先验框的序号，张量初始时所有值均置0，最匹配的先验框的序号位置置1。在2.4中用于获取对于每个预测框来说，找到与之iou值大于设定值的真实框的位置，更新无目标掩码张量。（这里在特别说下，2.3和2.4均用了该方法，其作用不同，2.3中是为了找到与预测框尺寸最匹配的先验框，2.4中是为了得到预测值后，与真实值进行比较，从而更新无目标掩码，即将交并比较大的位置不计入损失计算）
真实目标张量其实是非常稀疏的，对一张图片来说，若只有3个目标，则只有三个特征点位置的预测框属性有值，因此我们要找到这些特征点的位置、最匹配先验框序号和该点预测框的self.bbox_attrs属性，例如针对张量中一个具体目标来说，[x ,n,h,w,bbox_attrs ]，翻译过来就是说在第x张特征图上，在(h,w)特征点上，存在一个目标，该目标与序号n的先验框最匹配，该目标的真实框属性为bbox_attrs

方法源码：

def box_iou(self, box_a, box_b):
        """
        函数定义：计算一张图片上的所有真实框和先验框的IoU
        输入为：
         ----------
        box_a: tensor, shape=(gt_num, 4), xywh
        box_b: tensor, shape=(anchors_num, 4), xywh

        返回为：
        -------
        iou: tensor, shape=(gt_num, anchors_num, 1)
        """
        # 计算真实框的左上角和右下角坐标
        box_a_x1, box_a_x2 = box_a[:, 0] - box_a[:, 2] / 2, box_a[:, 0] + box_a[:, 2] / 2
        box_a_y1, box_a_y2 = box_a[:, 1] - box_a[:, 3] / 2, box_a[:, 1] + box_a[:, 3] / 2

        # 计算先验框的左上角和右下角坐标
        box_b_x1, box_b_x2 = box_b[:, 0] - box_b[:, 2] / 2, box_b[:, 0] + box_b[:, 2] / 2
        box_b_y1, box_b_y2 = box_b[:, 1] - box_b[:, 3] / 2, box_b[:, 1] + box_b[:, 3] / 2

        # 将真实框和先验框张量最后一维转化为左上角和右上角坐标形式，x1，y1, x2, y2
        _box_a = torch.zeros_like(box_a)
        _box_b = torch.zeros_like(box_b)
        _box_a[:, 0], _box_a[:, 1], _box_a[:, 2], _box_a[:, 3] = box_a_x1, box_a_y1, box_a_x2, box_a_y2
        _box_b[:, 0], _box_b[:, 1], _box_b[:, 2], _box_b[:, 3] = box_b_x1, box_b_y1, box_b_x2, box_b_y2

        # 计算真实框和先验框数量
        A = _box_a.size(0)
        B = _box_b.size(0)

        # 计算每个真实框与每个先验框的iou
        # 先将两个张量转化为相同维度，再进行计算
        # 计算所有交框的左上角和右下角坐标，再计算交框的面积
        intersect_mins = torch.max(_box_a[:, :2].unsqueeze(1).expand(A, B, 2), _box_b[:, :2].unsqueeze(0).expand(A, B, 2))  # [A, B, 2]
        intersect_maxs = torch.min(_box_a[:, 2:].unsqueeze(1).expand(A, B, 2), _box_b[:, 2:].unsqueeze(0).expand(A, B, 2))  # [A, B, 2]
        intersect_wh   = torch.max(intersect_maxs - intersect_mins, torch.zeros_like(intersect_maxs))  # [A, B, 2]
        intersect_area = intersect_wh[:, :, 0] * intersect_wh[:, :, 1]  # [A, B]

        # 计算并框的面积，再计算交并比
        box_a_area = ((_box_a[:, 2] - _box_a[:, 0]) * (_box_a[:, 3] - _box_a[:, 1])).unsqueeze(1).expand(A, B)  # [A, B]
        box_b_area = ((_box_b[:, 2] - _box_b[:, 0]) * (_box_b[:, 3] - _box_b[:, 1])).unsqueeze(0).expand(A, B)  # [A, B]
        iou = intersect_area / (box_a_area + box_b_area - intersect_area)  # [A, B]

        return iou

3.2.3 def get_target(self, l, targets, scale_anchors, in_h, in_w)

该方法主要是获取真实目标的张量（真实框张量）以及无目标掩码张量（box_loss_scale掩码在tiny模型中未使用），具体实现看源码吧，基本每行代码均有注释。

注意targets参数为列表而非张量，在计算时应选择对列表的操作方法
由于先验框只有宽高参数，并没有中心参数，因此在计算先验框和真实框的交并比前，应将二者的中心移至相同位置，这里的交并比并不是计算损失的交并比，而只是为了进行选择，对一个真实框来说，哪个先验框的尺寸更匹配
真实目标张量y_true的构建，b代表批次bs中的第b张图片，k代表这张图片上真实框最佳匹配的先验框序号，对tiny模型来说，先验框子表只有三个元素，则k取值为0,1,2，j和i代表真实目标的网格位置，tiny模型输出两路，13x13和26x26，因此j和i的取值对这两路分别为[0,12]，[0,25]，self.bbox代表真实框的属性，包括（5+classes）个元素，5代表xywhc，classes代表类别数量
关于无目标掩码张量noobj_mask的解释，在张量运算中常用到布尔掩码，通过布尔掩码进行数据筛选。对于无目标掩码张量，初始元素均置1，代表均无目标。然后根据真实框张量的情况，将有目标的位置元素置0

方法源码：

def get_target(self, l, targets, scale_anchors, in_h, in_w):
        """
        函数定义: 待填充
        :param l: l代表特征图序号，即选择第l个特征图
        :param targets: 代表真实框数据标签，包括批次，真实框数量，以及真实框坐标和类别信息，shape=[bs, gt_num, 5], 5->xywhc :注意，在后续调试过程，发现targets是一个列表，非张量
        :param scale_anchors: 在特征图尺度上的先验框列表, 即原anchors进行缩放
        :param in_h: 特征图高度
        :param in_w: 特征图宽度

        :return
        y_true:创建一个与模型输出形状相同的张量，用于存储真实的标签信息。在训练过程中，这个张量将用于计算损失函数，并反向传播到网络中。
        noobj_mask:创建一个张量，用于标记哪些先验框不包含物体（即负样本）。在训练过程中，这些负样本通常对损失函数的贡献较小，有助于平衡正负样本的影响
        box_loss_scale:用于调整不同大小的物体在损失函数中的权重，特别是用于让网络更加关注小目标。小目标在特征图中通常只占据少量像素，因此可能需要额外的权重来确保它们得到足够的关注。
        """
        # 获取批次大小，即图片的数量
        bs = len(targets)
        # 获取该特征图下的先验框数量
        anchors_num = len(self.anchors_mask[l])
        # 创建一个张量，标记哪些先验框不包含物体，初始化时全都不包含，即张量元素均为1
        noobj_mask = torch.ones(bs, anchors_num, in_h, in_w, requires_grad=False)
        # 创建一个张量，用于调整不同大小物体在损失函数中的权重，让网络更加关注小目标
        box_loss_scale = torch.zeros(bs, anchors_num, in_h, in_w, requires_grad=False)
        # 创建一个张量，与模型输出形状相同，存储真实的标签信息
        y_true = torch.zeros(bs, anchors_num, in_h, in_w, self.bbox_attrs, requires_grad=False)

        # 对批次内每张图片的每个真实框与预设所有先验框进行交并比计算，从而选出最大交并比的先验框，作为该真实框的最初尺寸分类
        for b in range(bs):
            if len(targets) == 0:
                continue
            batch_target = torch.zeros_like(targets[b])
            # 计算出原图中正样本在特征图上的中心点和宽高
            # 这里的target中的xywh可能为归一化坐标，故乘以特征图尺寸还原到特征图的像素坐标,因为后续先验框的尺寸为像素坐标，故两者统一
            batch_target[:, [0, 2]] = targets[b][:, [0, 2]] * in_w
            batch_target[:, [1, 3]] = targets[b][:, [1, 3]] * in_h
            # 提取真实框的类别序号
            batch_target[:, 4] = targets[b][:, 4]

            # 将真实框由列表转换为张量形式，中心坐标归0（先验框无中心坐标，所以两者都取0），即移至坐标原点,方便后续计算交并比,代入到box_iou进行计算
            # 真实框张量, shape = [gt_num, 4]
            gt_box = torch.FloatTensor(torch.cat((torch.zeros(batch_target.size(0), 2), batch_target[:, 2:4]), 1))
            # 先验框张量，shape = [anchors_num, 4]，这里的先验框数量取的是总数量，并非子列表的先验框数量,按照该类的先验框初始化情况，应为9个
            anchors_box = torch.FloatTensor(torch.cat((torch.zeros(len(scale_anchors), 2), torch.FloatTensor(scale_anchors)), 1))
            # 计算交并比
            # self.calculate_iou(gt_box, anchors_box) = [num_true_box, 9]每一个真实框和9个先验框的重合情况
            iou = self.box_iou(gt_box, anchors_box)  # [gt_num, anchors_num]
            # best_ns:
            # [每个真实框最大的重合度max_iou, 每一个真实框最重合的先验框的序号]
            best_ns = torch.argmax(iou, dim=-1)
            sort_ns = torch.argsort(iou, dim=-1, descending=True)

            # 检查最重合先验框的序号是否在先验框掩码里
            def check_in_anchors_mask(index, anchors_mask):
                for sub_anchors_mask in anchors_mask:
                    if index in sub_anchors_mask:
                        return True
                return False
            for t, best_n in enumerate(best_ns):
                if not check_in_anchors_mask(best_n, self.anchors_mask):
                    for index in sort_ns[t]:
                        if check_in_anchors_mask(index, self.anchors_mask):
                            best_n = index
                            break

                if best_n not in self.anchors_mask:
                    continue
                # 判断当前先验框是当前特征点的哪一个先验框，即确定当前特征点的最重合先验框的序号
                k = self.anchors_mask[l].index(best_n)
                # 获得真实框属于哪个网格点
                i = torch.floor(batch_target[t, 0]).long()
                j = torch.floor(batch_target[t, 1]).long()
                # 取出该真实框的标签种类，c从0开始，0对应第一个类别，依次类推
                c = batch_target[t, 4].long()
                # noobj_mask代表无目标的掩码，初始值均为1，若当前图片、当前先验框、当前特征点有目标，则置为0
                noobj_mask[b, k, j, i] = 0
                # 真实标签的张量y_true, shape=[bs, anchors_num, in_h, in_w, bbox_attrs]
                # 将真实框的中心点、宽高储存到当前特征点上
                y_true[b, k, j, i, 0] = batch_target[t, 0]
                y_true[b, k, j, i, 1] = batch_target[t, 1]
                y_true[b, k, j, i, 2] = batch_target[t, 2]
                y_true[b, k, j, i, 3] = batch_target[t, 3]
                # 将最后一维第5个元素置为1，即当前特征点一定含有目标，置信度为100%；将第c+5元素置为1，即当前特征点的该类别输出目标置1，其余类别输出目标仍为0
                y_true[b, k, j, i, 4] = 1
                y_true[b, k, j, i, c + 5] = 1

                # 损失函数中的不平衡：在目标检测任务中，尤其是当场景中同时存在大目标和小目标时，直接计算所有目标的损失可能会导致网络更关注大目标。
                # 这是因为大目标通常占据更多的像素，因此在损失计算中贡献更大。这可能导致网络对小目标的检测性能不佳。
                # 为了解决这个问题，可以通过给不同大小的目标分配不同的权重来调整损失函数。
                # 对于小目标，可以给予更高的权重，这样网络在训练时会更加关注小目标的预测误差，并尝试减小这些误差。
                # 此处计算存在归一化操作，归一化的宽在0-1之间，高在0-1之间，乘积也在
                box_loss_scale = 2-batch_target[t, 2] * batch_target[t, 3] / in_h / in_w
        return y_true, noobj_mask, box_loss_scale

3.2.4 def get_ignore(self, l, x, y, h, w, targets, scaled_anchors, in_h, in_w, noobj_mask)

该方法主要是获取预测结果的张量（预测框张量）以及更新无目标掩码张量，同样具体看源码。

先根据特征图尺寸生成网格，每个网格左上角坐标即为先验框的初始中心，因为对特征图来说，先验框有三个，则会在每个网格点生成三个预测框
先验框要先转化成特征图尺度的宽高，并将其张量形状转化为与输出y_true一致
对于每个特征点来说，每个先验框的中心和宽高有了之后，在加上特征图输出的xywh值，则得到预测框的张量
得到预测框后，需要对无目标掩码张量进行更新，更新的原则：对于该特征图的每个特征点的预测框，与该特征图的真实框进行交并比计算，若交并比大于设定值，则认为该特征点存在目标，则将对应位置的无目标掩码张量noobj_mask置0。注意，此处的预测框数量比较大，对于13x13特征图来说，预测框数量为13x13x3。

方法源码：

def get_ignore(self, l, x, y, h, w, targets, scaled_anchors, in_h, in_w, noobj_mask):
        """
        函数定义：在当前特征图l下，根据先验框以及模型输出值，求得预测框张量，并将每个预测框与真实框求交并比，将最大交并比的预测框（先验框）设为大概率有目标，故noobj_mask掩码对应位置置0，后续不参与损失计算
        :param l: l代表特征图序号，即选择第l个特征图
        :param x: 经过神经网络计算输出x，x代表先验框的调整量，从而得到预测框x
        :param y: 同x
        :param h: 同x，注意宽高调整量不能直接相加，因为已经经过对数处理，还原的话要在经指数处理
        :param w: 同h
        :param targets: 代表真实框数据标签，包括批次，真实框数量，以及真实框坐标和类别信息，shape=[bs, gt_num, 5], 5->xywhc
        :param scaled_anchors: 在特征图尺度上的先验框列表, 即原anchors进行缩放
        :param in_h: 特征图高度
        :param in_w: 特征图宽度
        :param noobj_mask: 无目标掩码
        :return:
        pred_boxes: 预测框张量，shape=[bs, len(anchors_mask[l]), in_h, in_w, 4]
        noobj_mask: 无目标掩码
        """
        # 获得批次大小
        bs = len(targets)
        # 生成网格，网格的左上角坐标即为先验框的中心
        grid_x = torch.linspace(0, in_w-1, in_w).repeat(in_h, 1).repeat(
            int(bs*len(self.anchors_mask[l])), 1, 1).view(x.shape).type_as(x)
        grid_y = torch.linspace(0, in_h-1, in_h).repeat(in_w, 1).t().repeat(
            int(bs*len(self.anchors_mask[l])), 1, 1).view(y.shape).type_as(x)

        # 生成先验框的在特征图尺度下的宽高，并转化为与输出相同的张量形状
        scaled_anchors_l = np.array(scaled_anchors)[self.anchors_mask[l]]  # scaled_anchors为元组列表，先转化为二维数组，并取出当前l掩码的宽高，一行代表一组宽高
        anchor_w = torch.Tensor(scaled_anchors_l).index_select(1, torch.LongTensor([0])).type_as(x)
        anchor_h = torch.Tensor(scaled_anchors_l).index_select(1, torch.LongTensor([1])).type_as(x)

        # 转换张量形状，用于后续计算
        anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h, in_w).view(w.shape)
        anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h, in_w).view(h.shape)

        # 计算调整之后的先验框的中心和宽高，也就是预测框的中心和宽高
        pred_boxes_x = torch.unsqueeze(grid_x + x, -1)
        pred_boxes_y = torch.unsqueeze(grid_y + y, -1)
        pred_boxes_w = torch.unsqueeze(torch.exp(w) + anchor_w, -1)
        pred_boxes_h = torch.unsqueeze(torch.exp(h) + anchor_h, -1)
        pred_boxes = torch.cat([pred_boxes_x, pred_boxes_y, pred_boxes_h, pred_boxes_w], -1)

        # 求得更新后的noobj_mask
        for b in range(bs):
            # 将第b批次的预测框张量转换形式，[len(anchors_mask[l]), in_h, in_w, 4] ->[anchors_sum_num, 4]
            pred_boxes_for_ignore = pred_boxes[b].view(-1, 4)  # B=len(anchors_mask[l])*in_h*in_w

            # 计算真实框张量，并且数量大于0，并转换为特征图尺度大小，shape=[gt_num, 4], A=4
            if len(targets[b]) > 0:
                batch_target = torch.zeros_like(targets[b])
                batch_target[:, [0, 2]] = targets[b][:, [0, 2]] * in_w
                batch_target[:, [1, 3]] = targets[b][:, [1, 3]] * in_h
                batch_target = batch_target[:, :4].type_as(x)

                # 计算所有预测框和所有真实框的交并比
                # 对于每一个预测框，求取最大交并比的真实框iou数值
                iou = self.box_iou(batch_target, pred_boxes_for_ignore)  # iou shape=[A, B]
                max_iou, _ = torch.max(iou, dim=0)  # max_iou为一维张量，大小为B
                # 将张量max_iou再转换为与第b批次的pred_boxes相同的形状
                max_iou = max_iou.view(pred_boxes[b].size()[:3])  # B -> [len(anchors_mask[l]), in_h, in_w]

                # 将max_iou的值与预设的ignore_threshord比较，大于预设交并比时为ture，形成一个布尔掩码。
                # 利用布尔掩码对noobj_mask对应位置进行操作，如果为真则置0
                # 该操作将预测框交并比较大的不予进入损失计算
                noobj_mask[b][max_iou > self.ignore_threshord] = 0

        return pred_boxes, noobj_mask

3.2.5 def forward(self, l, input, targets=None)

由于损失模块为继承nn.Module的子类，因此可以定义前向传播方法，因此可以通过实例调用该方法，完成损失计算。

关于参数的说明：l代表特征图序号，input代表图片经模型处理后的特征图输出，targets代表真实结果
先对input进行处理，形状shape如下

# shape: [bs, 3*(5+classes_num), 13, 13] -> [bs, 3, 13, 13, 5+classes_num]

# shape: [bs, 3*(5+classes_num), 26, 26] -> [bs, 3, 26, 26, 5+classes_num]
再调用get_target方法获得真实的结果y_true，y_true张量的形状与调整后的input张量形状相同。调用get_target可以获得第一次的无目标布尔掩码noobj_mask。
再调用get_ignore方法，获得预测的结果pred_boxes张量，该张量形状与y_true张量的形状相同。调用get_ignore方法可以获得更新后的无目标布尔掩码noobj_mask。后面根据y_true张量还要设定一个有目标的布尔掩码obj_mask
最后开始计算损失：
- 首先调用box_ciou方法计算y_true张量和pred_boxes张量的ciou。
- 计算回归损失，采用平均值。loss_loc = torch.mean((1 - ciou)[obj_mask])。通过有目标的布尔掩码obj_mask进行索引，我们只选择那些 obj_mask 为 True 的位置上的 (1 - ciou) 值。这确保了只有正样本的预测框对损失有贡献。
- 计算分类损失，采用二元交叉熵。loss_cls = torch.mean(self.BSELoss(pred_cls[obj_mask], y_true[..., 5:][obj_mask]))，同理使用布尔索引，确保了只有正样本的预测框分类对损失有贡献
- 合并回归损失和分类损失，并引入权重：loss += loss_loc * self.box_ratio + loss_cls * self.cls_ratio
- 置信度损失：采用二元交叉熵。loss_conf = torch.mean(self.BSELoss(conf, obj_mask.type_as(conf))[noobj_mask.bool() | obj_mask])。conf和obj_mask进行比较，由于布尔索引采用了noobj_mask与obj_mask的或逻辑运算，即考虑了正样本和负样本（这里只排除了中间样本，就是预测框和真实框iou大于设定值那部分样本）
- 最终损失值：loss += loss_conf * self.balance[l] * self.obj_ratio，对置信度损失引入权重值，大特征图权重大一些，让模型更关注小目标。

方法源码：

def forward(self, l, input, targets=None):
        """

        :param l: 第几个特征图序号
        :param input: 特征图输入，yolov4-tiny有两个特征图
                      l=0时，shape = [bs, 3*(5+classes_num), 13, 13]
                      l=1时，shape = [bs, 3*(5+classes_num), 26, 26]
        :param targets: 真实图输入，即真实标签输入情况，此处targets为列表非张量
                      shape = [bs, gt_num, 5]
        :return:
        loss: 输出损失值
        """
        # 获得批次大小，特征图宽和高的大小
        bs   = input.size(0)
        in_w = input.size(2)
        in_h = input.size(3)

        # 计算步长
        # 每一个特征点对应原来的图片上多少个像素点
        # 如果特征层为13x13的话，一个特征点就对应原来的图片上的32个像素点
        # 如果特征层为26x26的话，一个特征点就对应原来的图片上的16个像素点
        # stride_h = stride_w = 32、16
        stride_w = self.input_shape[1] / in_w
        stride_h = self.input_shape[0] / in_h

        # 计算scale_anchors，这个是相对特征图的先验框尺寸
        scale_anchors = [(anchor_w / stride_w, anchor_h / stride_h) for anchor_w, anchor_h in self.anchors]

        # 将输入的input张量进行转换，input张量一共有两个，对应两个特征图。将四维张量转换为五维。
        # shape: [bs, 3*(5+classes_num), 13, 13] -> [bs, 3, 13, 13, 5+classes_num]
        # shape: [bs, 3*(5+classes_num), 26, 26] -> [bs, 3, 26, 26, 5+classes_num]
        prediction = input.view(bs, len(self.anchors_mask[l]), self.bbox_attrs, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()

        # 先验框的中心调整参数，利用sigmoid函数将其映射至0-1之间，使其不会超出一个网格。后面再加上网格坐标，则得到预测框中心
        x = torch.sigmoid(prediction[..., 0])
        y = torch.sigmoid(prediction[..., 1])

        # 先验框的宽高调整参数
        w = prediction[..., 2]
        h = prediction[..., 3]

        # 由先验框得到的预测框的位置置信度
        conf = torch.sigmoid(prediction[..., 4])

        # 由先验框得到的预测框的种类置信度
        pred_cls = torch.sigmoid(prediction[..., 5:])

        # 获得网络应该得到的真实的预测结果，即目标值，y_true为真实框张量
        y_true, noobj_mask, box_loss_scale = self.get_target(l, targets, scale_anchors, in_h, in_w)

        # 将预测结果进行解码，获得网络根据先验框和计算结果得到的预测框, pred-boxes为预测框张量
        # 判断预测结果和真实值的重合程度，如果重合程度过大则忽略，因为这些特征点属于预测比较准确的特征点，作为负样本不合适
        pred_boxes, noobj_mask = self.get_ignore(l, x, y, h, w, targets, scale_anchors, in_h, in_w, noobj_mask)

        if self.cuda:
            y_true          = y_true.type_as(x)
            noobj_mask      = noobj_mask.type_as(x)
            box_loss_scale  = box_loss_scale.type_as(x)
        # --------------------------------------------------------------------------#
        #   box_loss_scale是真实框宽高的乘积，宽高均在0-1之间，因此乘积也在0-1之间。
        #   2-宽高的乘积代表真实框越大，比重越小，小框的比重更大。
        #   使用iou损失时，大中小目标的回归损失不存在比例失衡问题，故弃用
        # --------------------------------------------------------------------------#
        box_loss_scale = 2 - box_loss_scale

        # 计算损失
        loss = 0
        # 生成布尔掩码，真实框上为正样本，即相应位置置信度为1，即有目标的位置为True，其余不参与损失计算
        obj_mask = y_true[..., 4] == 1
        n = torch.sum(obj_mask)
        if n != 0:
            # 计算所有真实框和所有预测框的ciou
            ciou = self.box_ciou(y_true[..., :4], pred_boxes).type_as(x)
            # 计算回归损失，即定位回归损失值
            # 通过使用布尔索引，我们只选择那些 obj_mask 为 True 的位置上的 (1 - ciou) 值。这确保了只有正样本的预测框对损失有贡献。
            loss_loc = torch.mean((1 - ciou)[obj_mask])
            # 同理使用布尔索引，利用二元交叉熵损失函数计算分类损失
            loss_cls = torch.mean(self.BSELoss(pred_cls[obj_mask], y_true[..., 5:][obj_mask]))
            # 合并损失，引入权重
            loss += loss_loc * self.box_ratio + loss_cls * self.cls_ratio
        # 引入置信度损失
        # conf 与 obj_mask进行交叉熵损失计算，这计算了每个预测框的置信度损失。
        # 注意，这里实际上是将置信度预测与 obj_mask 作为目标值进行比较，这在YOLO中是一个常见的做法，因为YOLO将置信度解释为预测框内存在目标的概率。
        # noobj_mask用于标识哪些预测框没有与任何真实目标匹配（即负样本），由于在get_target和get_ignore函数中，对noobj_mask已更新，对部分位置不参与计算的预测框已经置0
        # noobj_mask为元素为0或1，noobj_mask.bool()将0转换为False，1转换为True
        # [noobj_mask.bool() | obj_mask]，这个逻辑或运算的结果是一个新的布尔数组，其中每个元素都是 noobj_mask 和 obj_mask 对应位置元素逻辑或的结果，形成新的布尔掩码
        # 因此置信度计算既考虑了正样本，也考虑负样本
        loss_conf = torch.mean(self.BSELoss(conf, obj_mask.type_as(conf))[noobj_mask.bool() | obj_mask])
        loss += loss_conf * self.balance[l] * self.obj_ratio
        return loss

3.3 其他方法（weights_init、get_lr_scheduler、set_optimizer_lr）

在损失模块中，最关键的就是YoloLoss类的建立，另外在这个模块中还编写了以下3个方法：weights_init：定义权重初始化函数；get_lr_scheduler：定义函数用于获取学习率调度器；set_optimizer_lr：定义函数用于设置优化器的学习率。

3.3.1 def weights_init(net, init_type=‘normal’, init_gain=0.02)

该方法用于初始化模型的所有参数，net是需要初始化的模型。该方法中又定义了一个def init_func(m)初始化方法，m是每个模块或层。

方法源码：

def weights_init(net, init_type='normal', init_gain=0.02):
    """

    :param net: 需要初始化的神经网络模型
    :param init_type: 初始化方法，默认为normal
    :param init_gain: 初始化增益，默认为0.02
    :return:
    对模型使用初始化方法，完成模型权重的初始化
    """
    def init_func(m):
        classname = m.__class__.__name__
        # 检测该层m是否包含权重属性'weight', 若包含则为TRUE
        # 检测该层m的类名中是否包含'Conv', 即是否为卷积层，若包含则返回索引，非-1，则逻辑判断为TRUE。若不包含则返回-1，逻辑判断为FALSE。
        if hasattr(m, 'weight') and classname.find('Conv') != -1:
            if   init_type == 'normal':
                torch.nn.init.normal_(m.weight.data, 0.0, init_gain)
            elif init_type == 'xavier':
                torch.nn.init.xavier_normal_(m.weight.data, gain=init_gain)
            elif init_type == 'kaiming':
                torch.nn.init.kaiming_normal_(m.weight.data, 0.0, 'fan_in')
            elif init_type == 'orthogonal':
                torch.nn.init.orthogonal_(m.weight.data, gain=init_gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
        # 检查该层是否是批量归一化层
        # 初始化权重正态分布，初始化偏置为常数0
        elif classname.find('BatchNorm2d') != -1:
            torch.nn.init.normal_(m.weight.data, 0.0, init_gain)
            torch.nn.init.constant_(m.bias.data, 0.0)
        print('initialize network with %s type' % init_type)
        net.apply(init_func)

3.3.2 def get_lr_scheduler(lr_delay_type, lr, min_lr, total_iters, warmup_iters_ratio=0.05, warmup_lr_ratio=0.1, no_aug_iter_ratio=0.05, step_num=10)

该方法主要是设置学习率调度器，学习率随着世代的变化进行变化，主要根据学习率变化类型有关，该方法设置了两种变化类型，一个是‘cos’余弦退火算法，一个‘step’步长递减算法。该方法调用时需要输入学习率变化类型，初始学习率，最小学习率以及总世代数，返回值为一个学习率关于世代的函数，注意这里返回的是一个函数。

方法源码：

def get_lr_scheduler(lr_delay_type, lr, min_lr, total_iters, warmup_iters_ratio=0.05, warmup_lr_ratio=0.1, no_aug_iter_ratio=0.05, step_num=10):
    """
    函数定义：函数 get_lr_scheduler，它用于获取学习率调度器。学习率调度器用于在训练神经网络时动态地调整学习率，以提高训练效率和模型性能。
    函数内部定义了两个内部函数 yolox_warm_cos_lr 和 step_lr，分别用于实现余弦退火策略和步长衰减策略。
    根据 lr_decay_type 的值，函数返回相应的学习率调度器函数 func
    :param lr_delay_type: 学习率衰减类型。如果为 "cos"，则使用余弦退火策略；否则，使用步长衰减策略。
    :param lr: 初始学习率。
    :param min_lr: 最小学习率，当使用余弦退火策略时，学习率将在这个值和初始学习率之间变化；当使用步长衰减策略时，学习率将逐渐衰减到这个值。
    :param total_iters: 总迭代次数，即整个训练过程的迭代次数
    :param warmup_iters_ratio: 预热迭代次数的比例，默认为0.05。预热阶段是在训练开始时逐渐增加学习率的过程，有助于模型更好地适应初始学习率。
    :param warmup_lr_ratio: 预热阶段开始时的学习率比例，默认为0.1。
    :param no_aug_iter_ratio: 不使用数据增强的迭代次数比例，默认为0.05。这部分迭代通常用于模型微调，此时可能不再使用数据增强。
    :param step_num: 步长衰减策略的步数，即学习率在每个步长后衰减一次。
    :return:
    func = functools.partial
    注： functools.partial是Python标准库中的一个函数，它的主要作用是部分应用一个函数，也就是固定函数的一部分参数，
    返回一个新的可调用对象。这个新的对象类似于原始函数，但其中的一些参数已经被预先设置，因此在后续调用中可以减少需要传递的参数数量
    """
    def yolox_warm_cos_lr(lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter, iters):
        if iters <= warmup_total_iters:
            lr = warmup_lr_start + (lr - warmup_lr_start) * pow(iters / warmup_total_iters, 2)
        elif iters >= total_iters - no_aug_iter:
            lr = min_lr
        else:
            lr = min_lr + 0.5 * (lr - min_lr) * (
                1.0 + math.cos(math.pi * ((iters - warmup_total_iters) / (total_iters - warmup_total_iters - no_aug_iter))))
        return lr

    def step_lr(lr, delay_rate, step_size, iters):
        if step_size < 1:
            raise ValueError("step_size must above 1.")
        else:
            n = iters // step_size
            lr = lr * delay_rate ** n
        return lr
    # 最后，根据 lr_decay_type 的值，函数返回相应的学习率调度器函数 func。如果 lr_decay_type 为 "cos"，
    # 则返回 yolox_warm_cos_lr 函数的部分应用（使用 functools.partial）；否则，返回 step_lr 函数的部分应用。
    if lr_delay_type == 'cos':
        warmup_toal_iters = min(max(total_iters * warmup_iters_ratio, 1), 3)
        warmup_lr_start = max(lr * warmup_lr_ratio, 1e-6)
        no_aug_iter = min(max(total_iters * no_aug_iter_ratio, 1), 15)
        func = functools.partial(yolox_warm_cos_lr, lr, min_lr, total_iters, warmup_toal_iters,
                                 warmup_lr_start, no_aug_iter)
    else:
        delay_rate = (min_lr / lr) ** (1 / (step_num - 1))
        step_size = total_iters / step_num
        func = functools.partial(step_lr, lr, delay_rate, step_size)

    return func

3.3.3 def set_optimizer_lr(optimizer, lr_scheduler_func, epoch)

该方法比较简单，主要用于设置pytorch的优化器的学习率调度器

方法源码：

def set_optimizer_lr(optimizer, lr_scheduler_func, epoch):
    """
    函数定义定义了一个函数 set_optimizer_lr，它用于设置 PyTorch 优化器的学习率。这个函数接受三个参数：，lr_scheduler_func（学习率调度函数），和 epoch（当前的训练周期数）
    :param optimizer: optimizer（优化器对象）
    :param lr_scheduler_func: lr_scheduler_func（学习率调度函数）
    :param epoch: epoch（当前的训练周期数）
    """
    lr = lr_scheduler_func(epoch)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

四. 工具模块

源码中的工具文件夹中有多个模块，这里写把训练所需的列出来，包括utils.py、callbacks.py、dataloader.py和utils_oneepoch.py

4.1 utils.py常见工具

utils.py程序中存放了常见的几个方法：

4.1.1 def cvtColor(image)

该方法输入为图片（格式为PIL类），如果图片是RGB图像，则返回图像，如果图像是灰度图像，则转换成RGB格式再返回。

方法源码：

# ---------------------------------------#
# 将图像转换成RGB图像，防止灰度图在预测时报错
# 代码仅支持RGB图像的预测，所有其他类型的图像都会转化为RGB
# ---------------------------------------#

def cvtColor(image):
    if isinstance(image, Image.Image):
        image_array = np.array(image)
        if len(np.shape(image_array)) == 3 and np.shape(image_array)[2] == 3:
            return image
        else:
            image = image.convert('RGB')
            return image

4.1.2 def get_classes(classes_path)

该方法输入为分类文件的路径，返回类名列表和类的数量。

方法源码：

# -------------------------------------- #
# 读取分类文件，获得类名和数量
# -------------------------------------- #
def get_classes(classes_path):
    with open(classes_path, encoding='utf-8') as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names, len(class_names)

4.1.3 def get_anchors(anchors_path)

该方法输入为先验框文件的路径，返回先验框数组和先验框数量

方法源码：

# -------------------------------------- #
# 读取先验框文件，获得先验框数组和先验框数量
# -------------------------------------- #
def get_anchors(anchors_path):
    with open(anchors_path, encoding='utf-8') as f:
        anchors = f.readline()
    anchors = [float(x) for x in anchors.split(',')]
    anchors = np.array(anchors).reshape(-1, 2)
    return anchors, len(anchors)

4.1.4 def get_lr(optimizer)

该方法是获得优化器中的学习率

方法源码：

# ---------------------------------------#
# 获得优化器中学习率
# ---------------------------------------#
def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']
# ---------------------------------

4.1.5 def seed_everything(seed=11)和def worker_init_fn(worker_id, rank, seed)

这俩方法主要和种子有关，我也没理解太好

4.1.6 def preprocess_input(image)

该方法对图片中的像素进行归一化操作。

方法源码：

def preprocess_input(image):
    image /= 255.0
    return image

4.1.7 def show_config(**kwargs)

该方法主要是用于在训练时显示所有初始化参数，便于训练初期查看。

方法源码：

# ---------------------------------------#
# 打印出训练参数
# ---------------------------------------#
def show_config(**kwargs):
    print('Configurations:')
    print('-'*70)
    print('|%25s | %40s|' % ('keys', 'values'))
    print('-'*70)
    for key, value in kwargs.items():
        print('|%25s | %40s|' % (str(key), str(value)))
    print(('-'*70))

4.2 callbacks.py

callbacks.py程序中存放了记录损失历史的模块class LossHistory()，该模块内定义了两个方法def append_loss(self, epoch, loss, val_loss)和def loss_plot(self)。这部分比较简单，直接贴上源码看注释吧。。。

类的源码：

class LossHistory():
    def __init__(self, log_dir, model, input_shape):
        self.log_dir = log_dir
        self.losses = []
        self.val_losses = []

        os.makedirs(self.log_dir)
        self.writer = SummaryWriter(self.log_dir)
        try:
            dummy_input = torch.randn(2, 3, input_shape[0], input_shape[1])
            self.writer.add_graph(model, dummy_input)
        except:
            pass

    # 将当前世代epoch的训练损失和验证损失写入到列表中和文件中，并绘制折线图
    def append_loss(self, epoch, loss, val_loss):
        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir)

        self.losses.append(loss)
        self.val_losses.append(val_loss)

        with open(os.path.join(self.log_dir, 'epoch_loss.txt'), 'a') as f:
            f.write(str(loss))
            f.write('\n')
        with open(os.path.join(self.log_dir, 'epoch_val_loss.txt'), 'a') as f:
            f.write(str(val_loss))
            f.write('\n')
        self.writer.add_scalar('loss', loss, epoch)
        self.writer.add_scalar('val_loss', val_loss, epoch)
        self.loss_plot()

    # loss_plot 的方法，用于绘制训练损失（train loss）和验证损失（val loss）的曲线图
    def loss_plot(self):
        iters = range(len(self.losses))

        plt.figure()
        plt.plot(iters, self.losses, 'red', linewidth=2, label='train loss')
        plt.plot(iters, self.val_losses, 'coral', linewidth=2, label='val loss')
        try:
            if len(self.losses) < 25:
                num = 5
            else:
                num = 15

            plt.plot(iters, scipy.signal.savgol_filter(self.losses, num, 3), 'green', linestyle='--', linewidth=2, label='smooth train loss')
            plt.plot(iters, scipy.signal.savgol_filter(self.val_losses, num, 3), '#8B4513', linestyle='--', linewidth=2, label='smooth val loss')
        except:
            pass

        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend(loc="upper right")

        plt.savefig(os.path.join(self.log_dir, "epoch_loss.png"))

        plt.cla()
        plt.close("all")

4.3 dataloader.py数据加载器（含数据增强）

dataloader.py程序中存放了数据加载的模块class YoloDataset(Dataset)，该模块继承了pytorch的Dataset类。数据加载器主要是对数据集图片进行预处理（或叫增强处理），针对训练集和验证集的操作不同。

针对训练集中的每一张图片，先第一步判断是否进行mosaic增强，并将增强后的图片和真实框返回。再第二部判断是否进行mixup增强，并将增强后的图片和真实框返回。如果第一步判断不进行的话，则执行常规的图像随机数据处理。
常规的图像随机数据处理包括图像缩放、空白处填充。函数里对训练集图片的处理和验证集图片的处理不同：训练集图片缩放随机、宽高比也被扭曲，训练图片的真实框也同理；验证集缩放至与input_shape大小一致，宽高比也保持不变，验证图片的真实框也同理。
这部分代码撸的时候，因为train.py在设置时mosaic和mixup两种模式都关闭了，所以相应的方法也没写，看源码吧，同志们

程序源码：

class YoloDataset(Dataset):
    def __init__(self, annotation_lines, input_shape, num_classes, epoch_length,
                        mosaic, mixup, mosaic_prob, mixup_prob, train, special_aug_ratio=0.7):
        super().__init__()
        self.annotation_lines = annotation_lines
        self.input_shape      = input_shape
        self.num_classes      = num_classes
        self.epoch_length     = epoch_length
        self.mosaic           = mosaic
        self.mosaic_prob      = mosaic_prob
        self.mixup            = mixup
        self.mixup_prob       = mixup_prob
        self.train            = train
        self.special_aug_ration = special_aug_ratio

        self.epoch_now = -1
        self.length = len(self.annotation_lines)

    # 当你创建了一个这个类的实例，并尝试使用 len() 函数来获取其长度时，
    # __len__ 方法会被调用，并返回 self.length 的值
    def __len__(self):
        return self.length

    # __getitem__：这是一个特殊方法，它允许类的实例使用方括号索引操作，比如dataset[index]
    def __getitem__(self, index):
        # 取模运算来确保index在0到self.length - 1的范围内,
        # 当需要多次遍历整个数据集时, 通常用于实现数据集的循环加载
        index = index % self.length

        # ----------------------------------------#
        # 训练时进行数据的随机增强
        # 验证时不进行数据的随机增强
        # ----------------------------------------#
        # 判断是否进行mosaic增强
        if self.mosaic and self.rand() < self.mosaic_prob and self.epoch_now < self.epoch_length * self.special_aug_ration:
            # 从`self.annotation_lines`中随机选择3个样本
            lines = sample(self.annotation_lines, 3)
            # 将当前索引`index`对应的样本添加到`lines`中
            lines.append(self.annotation_lines[index])
            # 打乱`lines`的顺序
            shuffle(lines)
            # 使用函数进行Mosaic数据增强，得到增强后的图像`image`和对应的边界框`box`
            image, box = self.get_random_data_with_Mosaic(lines, self.input_shape)

            # 如果进行mosaic增强，再判断是否进行mixup增强
            if self.mixup and self.rand() < self.mixup_prob:
                # 从'annotation_lines'中随机选择1个样本
                lines = sample(self.annotation_lines, 1)
                # 由于lines中只有1个元素，lines[0]则是这个样本本身，代入常规随机处理函数，返回新的图片和真实框
                image_2, box_2 = self.get_random_data(lines[0], self.input_shape, random=self.train)
                image, box = self.get_random_data_with_Mixup(image, box, image_2, box_2)
        # 如果不进行mosaic增强，则对数据进行常规随机处理
        else:
            image, box = self.get_random_data(self.annotation_lines[index], self.input_shape, random=self.train)

        image = np.transpose(preprocess_input(np.array(image, dtype=np.float32)), (2, 0, 1))
        box   = np.array(box, dtype=np.float32)
        if len(box) != 0:
            box[:, [0, 2]] = box[:, [0, 2]] / self.input_shape[1]
            box[:, [1, 3]] = box[:, [1, 3]] / self.input_shape[0]

            box[:, 2:4] = box[:, 2:4] - box[:, 0:2]
            box[:, 0:2] = box[:, 0:2] + box[:, 2:4] / 2
        return image, box

    # 用于生成一个在闭区间 [a, b] 内的随机浮点数
    def rand(self, a=0.0, b=1.0):
        return np.random.rand() * (b -a) +a

    # ------------------------------#
    # 常规随机处理函数，包括图像缩放、空白处填充
    # 函数里对训练集图片的处理和验证集图片的处理不同
    # 训练集图片缩放随机、宽高比也被扭曲，训练图片的真实框也同理；
    # 验证集缩放至与input_shape大小一致，宽高比也保持不变，验证图片的真实框也同理。
    def get_random_data(self, annotation_line, input_shape, jitter=0.3, hue=0.1, sat=0.7, val=0.4, random=True):
        line = annotation_line.split()
        # ---------------------------#
        # 读取图像，并转换成RGB图像
        # ---------------------------#
        image = Image.open(line[0])
        image = cvtColor(image)
        # ---------------------------#
        # 获得图像的高宽与目标高宽
        # ---------------------------#
        iw, ih = image.size
        h, w   = input_shape
        # ---------------------------#
        # 获得预测框，将所有预测框转化成二维数组
        # ---------------------------#
        box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])

        # random = self.train，
        # 即验证集不参与训练，图片则直接进行以下处理（图片缩放和真实框缩放）
        if not random:
            # 对原始图像进行处理
            # 取原始图像较长的边作为缩放比例尺
            scale = min(w/iw, h/ih)
            # 对原始图像的宽高进行缩放
            # 较长边缩放与input_shape相同长度，较短边缩放比input_shape要短，存在空余部分
            nw = int(iw * scale)
            nh = int(ih * scale)
            dx = (w-nw) // 2
            dy = (h-nh) // 2

            # 将图像空余的部分加上灰条
            image = image.resize((nw, nh), Image.BICUBIC)
            new_image = Image.new('RGB', (w, h), (128, 128, 128))
            new_image.paste(image, (dx, dy))
            image_data = np.array(new_image, np.float32)

            # 对原始图像的真实框进行处理,真实框采用左上右下坐标形式
            if len(box) > 0:
                np.random.shuffle(box)
                # 对真实框进行缩放，并且加上位置偏置
                box[:, [0, 2]] = box[:, [0, 2]] * nw / iw + dx
                box[:, [1, 3]] = box[:, [1, 3]] * nh / ih + dy
                # 通过布尔数组的方式，对真实框左上右下坐标进行处理，使其在图像范围内
                # 对真实框左上坐标进行不小于0处理
                # 对真实框右下坐标进行不大于缩放后宽高处理
                box[:, 0:2][box[:, 0:2] < 0] = 0
                box[:, 2][box[:, 2] > w] = w
                box[:, 3][box[:, 3] > h] = h
                box_w = box[:, 2] - box[:, 0]
                box_h = box[:, 3] - box[:, 1]
                # 宽度列和高度列生产布尔数组
                # 对两个布尔数组进行逻辑与计算，生成新的布尔数据
                # 新的布尔数组用于选取宽高大于1的真实框作为有效真实框
                box = box[np.logical_and(box_w > 1, box_h > 1)]

            return image_data, box

        # ---------------------------------------------#
        # random = self.train，
        # 训练集图片进行处理，对图像进行缩放并且进行长和宽的扭曲
        # ---------------------------------------------#
        # 定义新的宽高比，如果new_ar>1，则宽长，反之则高长
        new_ar = iw / ih * self.rand(1 - jitter, 1 + jitter) / self.rand(1 - jitter, 1 + jitter)
        scale  = self.rand(0.25, 2.0)
        if new_ar < 1:
            nh = int(ih * scale)
            nw = int(nh * new_ar)
        else:
            nw = int(iw * scale)
            nh = int(nw / new_ar)
        # 将图像空余的部分加上灰条，前提是有空余部分，因为缩放比例随机、宽高比也被调整
        image = image.resize((nw, nh), Image.BICUBIC)
        dx = int(self.rand(0, w - nw))
        dy = int(self.rand(0, h - nh))
        new_image = Image.new('RGB', (w, h), (128, 128, 128))
        new_image.paste(image, (dx, dy))
        image = new_image

        # ---------------------------------------------#
        # 翻转图像,翻转概率为50%
        # ---------------------------------------------#
        flip = self.rand() < 0.5
        if flip:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
        image_data = np.array(image, np.uint8)
        # ------------------------------------#
        # 对图像进行色域变换
        # 计算色域变换的参数
        # ------------------------------------#
        # 生成色调、饱和度、亮度的扰动值
        r = np.random.uniform(-1, 1, 3) * [hue, sat, val] + 1
        # 将图像转到HSV上
        hue, sat, val = cv2.split(cv2.cvtColor(image_data, cv2.COLOR_RGB2HSV))
        dtype = image_data.dtype
        # --------------------#
        # 应用变换
        # --------------------#
        # 用arange函数来创建一个一维数组 x。这个数组从 0 开始，到 256（不包括 256）结束，步长为 1
        # 数据类型是 np.uint8（无符号8位整数），那么 x 也将是一个数据类型为 np.uint8 的一维数组，包含了从 0 到 255 的整数
        x = np.arange(0, 256, dtype=dtype)
        # 计算查找表LUT
        lut_hue = ((x * r[0]) % 180).astype(dtype)
        lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
        lut_val = np.clip(x * r[2], 0, 255).astype(dtype)
        # 应用查找表
        image_data = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
        image_data = cv2.cvtColor(image_data, cv2.COLOR_HSV2RGB)
        # -------------------------------------#
        # 对真实框进行调整
        # -------------------------------------#
        if len(box) > 0:
            np.random.shuffle(box)
            box[:, [0, 2]] = box[:, [0, 2]] * nw / iw + dx
            box[:, [1, 3]] = box[:, [1, 3]] * nh / ih + dy
            if flip:
                box[:, [0, 2]] = w - box[:, [2, 0]]
            box[:, 0:2][box[:, 0:2] < 0] = 0
            box[:, 2][box[:, 2] > w] = w
            box[:, 3][box[:, 3] > h] = h
            box_w = box[:, 2] - box[:, 0]
            box_h = box[:, 3] - box[:, 1]
            box = box[np.logical_and(box_w > 1, box_h > 1)]

        return image_data, box

    # def merge_bboxes(self, bboxes, cutx, cuty):
    #
    #     return merge_bbox
    # 定义mosaic数据增强函数
    # def get_random_data_with_Mosaic(self, annotation_line, input_shape, jitter=0.3, hue=.1, sat=0.7, val=0.4):
    #
    #     return new_image, new_boxes
    # ------------------------------------#
    # 定义mixup增强函数
    # ------------------------------------#
    # def get_random_data_with_Mixup(self, image_1, box_1, image_2, box_2):
    #
    #     return new_image, new_boxes

def yolo_dataset_collate(batch):
    images = []
    bboxes = []
    for img, box in batch:
        images.append(img)
        bboxes.append(box)
    images = torch.from_numpy(np.array(images)).type(torch.FloatTensor)
    bboxes = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in bboxes]
    return images, bboxes

4.4 utils_oneepoch.py一个世代内训练

utils_oneepoch.py程序中只定义了一个方法def one_epoch(model_train, model, yolo_loss, loss_history, optimizer, epoch, epoch_step, epoch_val_step, gen, gen_val, Epoch, Cuda, fp16, scaler, save_period, save_dir, local_rand=0)。这个方法的名字我修改过了，与源码不同。其内容写的是在一个世代内进行训练的方法。根据神经网络训练的步骤，一般包括以下几个步骤：**1.模型切换至训练模式 2.清理之前的梯度 3.前向传播，得到模型的输出 4.根据模型的输出和真实结果，进行损失计算 5.反向传播，计算梯度 6.使用优化器更新模型参数。**在一个世代中，先进行训练集训练，再进行验证集验证。通过进行训练模式和验证模式的切换，以及数据集加载的不同来实现各自的完整步骤。主要的区别是，在进行验证集验证时，只进行到第4步损失计算，不进行反向传播和权重的优化。

程序源码：

# 定义一个世代的训练函数
def one_epoch(model_train, model, yolo_loss, loss_history,
              optimizer, epoch, epoch_step, epoch_val_step, gen, gen_val,
              Epoch, Cuda, fp16, scaler, save_period, save_dir, local_rand=0):
    """

    :param model_train: 表示模型所处状态，True:训练状态，False：验证/推理状态
    :param model: 模型本身
    :param yolo_loss: 损失模块
    :param loss_history: 损失记录模块
    :param eval_callback: 评估模块
    :param optimizer: 优化器模块
    :param epoch:   当前处于训练/验证第几世代
    :param epoch_step: 每个训练世代的步长
    :param epoch_val_step: 每个验证世代的步长
    :param gen: 训练集数据加载器
    :param gen_val: 验证集数据加载器
    :param Epoch: 总的训练世代，取值为UnFreeze_Epoch
    :param Cuda: 是否启用Cuda：bool
    :param fp16: 是否启动半精度计算：bool
    :param scaler: 同fp16状态相关的参数
    :param save_period: 每多少世代保存一次
    :param save_dir: 保存的目录
    :param local_rand:
    :return:
    """
    # 初始化损失值
    loss = 0
    val_loss = 0

    # ---------------------------#
    # 开始进行训练集训练，一般包括以下几步：
    # 1.模型切换至训练模式
    # 2.清理之前的梯度
    # 3.前向传播，得到模型的输出
    # 4.根据模型的输出和真实结果，进行损失计算
    # 5.反向传播，计算梯度
    # 6.使用优化器更新模型参数
    # ---------------------------#

    if local_rand == 0:
        print('开始第', epoch + 1, '世代训练')
        # 设置训练进度显示，利用tqdm进度条库建立一个pbar进度条对象
        # tqdm进度条在循环中通常与update()方法一起使用，以便在每次迭代时更新进度
        pbar = tqdm(total=epoch_step, desc=f'Epoch: {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3)
    # 模型切换至训练模式
    model_train.train()
    for iteration, batch in enumerate(gen):
        if iteration >= epoch_step:
            break
        # 从当前批次batch中提取图像（images）和目标（targets）
        images, targets = batch[0], batch[1]
        with torch.no_grad():
            if Cuda:
                images = images.cuda(local_rand)
                targets = [ann.cuda(local_rand) for ann in targets]
        # ----------------------#
        # 清零梯度
        # ----------------------#
        optimizer.zero_grad()
        # 如果非半精度计算的话，开始进行前向传播、计算损失、反向传播的操作
        if not fp16:
            # ------------------#
            # 前向传播
            # ------------------#
            outputs = model_train(images)

            loss_value_all = 0
            # ------------------#
            # 计算损失
            # ------------------#
            for l in range(len(outputs)):
                loss_item = yolo_loss(l, outputs[l], targets)
                loss_value_all += loss_item
            loss_value = loss_value_all

            # -----------------#
            # 反向传播
            # loss_value通常是一个标量（scalar）张量，它代表了模型在当前批次数据上的损失。
            # 调用backward()方法后，PyTorch会自动计算loss_value关于模型中所有可训练参数（即需要优化的参数）的梯度，
            # 并将这些梯度存储在对应参数的.grad属性中
            # optimizer.step()会根据存储在模型参数.grad属性中的梯度来更新这些参数。
            # 具体来说，它会根据优化器内部设定的学习率和其他超参数来调整每个参数的值，以减小损失函数。
            # 在optimizer.step()被调用之前，你需要先调用optimizer.zero_grad()来清除之前计算得到的梯度（如果有的话），因为PyTorch会累积梯度
            # -----------------#
            loss_value.backward()
            optimizer.step()
        else:
            from torch.cuda.amp import autocast
            with autocast():
                # 前向传播
                outputs = model_train(images)
                # 计算损失
                loss_value_all = 0
                for l in range(len(outputs)):
                    loss_item = yolo_loss(l, outputs[l], targets)
                    loss_value_all += loss_item
                loss_value = loss_value_all
            # 反向传播
            scaler.scale(loss_value).backward()
            scaler.step(optimizer)
            scaler.update()

        # 计算截止当前步长的损失总值
        loss += loss_value.item()

        if local_rand == 0:
            pbar.set_postfix(**{'loss': loss / (iteration + 1),
                                'lr': get_lr(optimizer)})
            pbar.update(1)

    if local_rand == 0:
        pbar.close()
        print('结束第', epoch + 1, '世代训练')
    # 开始进行验证集推理
        print('开始第', epoch + 1, '世代验证')
        pbar = tqdm(total=epoch_val_step, desc=f'Epoch: {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3)
    # 模型切换至验证模式
    model_train.eval()
    for iteration, batch in enumerate(gen_val):
        if iteration >= epoch_val_step:
            break

        images, targets = batch[0], batch[1]
        with torch.no_grad():
            if Cuda:
                images = images.cuda(local_rand)
                targets = [ann.cuda(local_rand) for ann in targets]
            # 清理梯度
            optimizer.zero_grad()
            # 前向传播
            outputs = model_train(images)
            # 计算损失
            loss_value_all = 0
            for l in range(len(outputs)):
                loss_item = yolo_loss(l, outputs[l], targets)
                loss_value_all += loss_item
            loss_value = loss_value_all
        val_loss += loss_value.item()

        if local_rand == 0:
            pbar.set_postfix(**{'val_loss': val_loss / (iteration + 1)})
            pbar.update(1)

    if local_rand == 0:
        pbar.close()
        print('结束第', epoch + 1, '世代验证')

        loss_history.append_loss(epoch + 1, loss / epoch_step, val_loss / epoch_val_step)
        # eval_callback.on_epoch_end(epoch + 1, model_train)
        print('Epoch:' + str(epoch + 1) + '/' + str(Epoch))
        print('Total Loss: %.3f || Val loss: %.3f' %(loss / epoch_step, val_loss / epoch_val_step))

        # ----------------------------------#
        # 保存权值
        # ----------------------------------#
        if (epoch + 1) % save_period == 0 or epoch + 1 ==Epoch:
            torch.save(model.state_dict(), os.path.join(save_dir, "ep%03d-loss%.3f-val_loss%.3f.pth" % (epoch + 1, loss / epoch_step, val_loss / epoch_val_step)))

        if len(loss_history.val_losses) <= 1 or (val_loss / epoch_val_step) <= min(loss_history.val_losses):
            print('Save best model to best_epoch_weights.pth')
            torch.save(model.state_dict(), os.path.join(save_dir, "best_epoch_weights.pth"))

        torch.save(model.state_dict(), os.path.join(save_dir, "last_epoch_weights.pth"))

五. 训练模块

5.1 参数设置

需要设置的参数比较多，这里只讲几个重点参数：

classes_path = '../model_data/voc_classes.txt'# 训练前一定要修改classes_path，使其对应自己的数据集
anchors_path = '../model_data/yolo_anchors.txt'# anchors_path 代表先验框对应的txt文件，一般不修改
model_path = '../model_data/yolov4_tiny_weights_coco.pth'#整个模型权重的路径，已经训练好的权重，直接用
pretrained = False是否使用主干网络的预训练权重，由于model_path使用的是全模型的权重，如果设置，则pretrained值无意义，如果model_path无设置，此处为True时则加载主干的预训练权重，为False则不加载，模型从0开始训练。
Freeze_Epoch = 50 Freeze_batch_size = 32 UnFreeze_Epoch = 300 UnFreeze_batch_size = 16 Freeze_Train = False。这几个参数是设置冻结训练使用的，冻结训练冻结的是主干网络的权重，也就是在冻结世代时，主干网络的权重不发生变化，只对其他权重进行调整。这个与是否使用主干预训练权重不冲突，两者概念不同。
train_annotation_path = '2007_train.txt' val_annotation_path = '2007_val.txt' 训练集和验证集的索引文件。
另外参数包括初始化学习率、最小学习率（这两个参数还会根据后面优化器的选择进一步变化）、学习率衰减方式、优化器类型等等参数，具体查看代码。

5.2 模型建立及参数初始化、损失模块初始化、数据集初始化

5.2.1 初始化模型参数，分三种情况：

pretrained = False，model_path = '': 既不加载主干预训练权重，也不加载整个模型的预训练权重，进行模型权重初始化（调用weights_init(model)方法，权重基本都是正态分布初始化），模型从0开始训练。
pretrained = True，model_path = ''：加载主干预训练权重，加载过程在建立模型实例过程中已加载（def darknet53_tiny方法中加载）
pretrained = True/False，model_path = 'XXX.pth'：不管True或False，主干预训练权重要么未加载，要么加载后也被整个模型的权重覆盖，因此此时模型的权重是整个模型的预训练权重。
再加载整个模型的预训练权重时，可以看到权重参数都是以字典形式存储，其中键是参数的名称，值是参数的张量。在加载时对预训练权重的参数逐一检查，如果键在当前模型参数中，且值的形状与当前模型参数一致，则加载。

5.2.2 损失模块初始化

调用损失类YoloLoss构建损失实例，并建立LossHistory实例对训练集损失和验证集损失进行记录。

5.2.3 数据集初始化

加载训练集和验证集的索引文件，逐行读取并将其转化成字符串列表。在数据加载器get_random_data方法中，对每行字符串进行操作，例如第一行："VOCdevkit/VOC2007/JPEGImages/000005.jpg 263,211,324,339,8 165,264,253,372,8 241,194,295,299,8"，首先对该行字符串按照空格进行分割，形成列表["VOCdevkit/VOC2007/JPEGImages/000005.jpg", "263,211,324,339,8", "165,264,253,372,8", "241,194,295,299,8"]，然后加载列表第一个元素，也就是图片对象。后面的元素均为box的参数，可以看到三个真实框，以及每个真实框的左上右下坐标和类别值。
在加载数据集的过程中，会根据数据集的数量进行提前判断，避免过小数据集进行训练。

5.3 开始训练

5.3.1 学习率初始化+优化器初始化

前面设置的初始学习率参数和最小学习率参数，根据当前batch_size的大小，以及优化器类型自适应调整
选择优化器类型，并将模型的参数按照不同的类型进行分组，便于针对不同类型参数设置不同的优化策略。其中偏置属性’bias’归为一组，批量归一化模块中的权重属性’weight’归为一组，剩余其他模块中包括权重属性’weight’的归为一组。

5.3.2 数据加载器进行数据加载

利用YoloDataset类建立训练集train_dataset和验证集实例val_dataset，该实例可以通过使用方括号索引操作，比如train_ataset[index],从而从数据集中选择图像，完成数据加载和增强处理，最后返回处理后的图像和真实框。
再调用torch中的数据加载器标准类DataLoader构建待训数据gen和gen_val

5.3.3 构建世代训练循环

世代训练循环for epoch in range(Init_Epoch, UnFreeze_Epoch)，开始进行逐代循环训练。
在世代训练循环中，首先判断当前世代是否在冻结或非冻结世代，如果进入非冻结世代，则需要对部分参数进行调整。并将解冻标识置UnFreeze_flag置为true，以免每个非冻结世代都要重置一次部分参数（参数包括学习率lr_scheduler_func、步长epoch_step/epoch_step_val和数据集gen/gen_val等，因为冻结世代和非冻结世代的batch_size不同，因此学习率也不同。而batch_size改变，数据集也要重新构建成新的gen和gen_val）。
然后不管是冻结世代和非冻结世代，在循环中调用set_optimizer_lr，根据当前世代设置学习率调度器。最后调用one_epoch（在循环中将调用之前写好的一个世代内的训练方法one_epoch，一个世代内是按照步长进行迭代的，将所有数据集完全训练一次。）

5.4 训练源码

这部分源码调用了之前所有构建的模块/类/方法/函数等内容，从我对代码的理解算是主脉络，从这里可以清楚的看到整个训练的全部过程以及前面构建内容的目的。

代码因为注释比较多，~~注释远远多余代码，我可太厉害了~~，我就不贴了。附上链接可自行下载。链接是直接到train.py。

Yolov4-tiny-clone/client at main · iton1ght/Yolov4-tiny-clone

六. 总结

没什么好总结的，请把前面的内容多看几遍吧，一行一行的撸代码。

欢迎提问交流~

itoN1ght

关注

47
点赞
踩
39

收藏

觉得还不错? 一键收藏
0
评论
YOLOv4-tiny（一）模型构建与训练实现

YOLOv4-tiny模型是在YOLOv4模型基础上优化裁剪的轻量化模型，模型参数量只有600万（YOLOv4有6000万），因此该模型速度非常快，准确率也还可以，对于想入门学习的小白我个人觉得还是比较合适的。我个人的学习路线是：1.下载模型源码 2.部署运行环境 3.进行推理测试 4.关键步骤：对照模型架构使劲撸源码。我是纯外行，所以感觉先调通模型会比较有成就感，虽然代码看起来很吃力，但是当视频中我的脑袋被检测到的时候还是很开心的。借着这篇文章算是一个学习记录，也希望能帮助有需要的人。
复制链接

扫一扫