YOLOv10（2）：网络结构及其检测模型代码部分阅读

赛先生.AI

已于 2024-05-31 09:19:54 修改

阅读量7.9k

点赞数 37

分类专栏： YOLOv10 计算机视觉文章标签： YOLOv10 计算机视觉

于 2024-05-31 09:19:22 首次发布

本文链接：https://blog.csdn.net/tecsai/article/details/139339772

版权

计算机视觉同时被 2 个专栏收录

13 篇文章

订阅专栏

YOLOv10

5 篇文章

订阅专栏

YOLOv10（1）：初探，训练自己的数据-CSDN博客

（3）PSA(partial self-attention)

3. 代码解读

（1）BaseModel(nn.Module)

（2）DetectionModel(BaseModel)

（3）YOLOv10DetectionModel(DetectionModel)

（4）拓展讲解：parse_model

1. 写在前面

YOLOv10的代码结构和YOLOv8很相似，说是基于YOLOv8魔改的也不为过。话说回来，也算是魔改的非常成功的了。从代码工程上看，也是通过类一层层继承，获得一个完整的模型网络。

2. 局部模块

YOLOv10使用了YOLOv8同样的C2f结构，在此就不多做解释了，在我的其他文章中均有讲解。

本次着重在两个模块上做一下讲解，分别是SCDown、C2fCIB以及PSA。

相关的代码在./ultralytics/nn/modules/block.py文件中。

（1）SCDown

SCDown理解起来相对简单，只由两个卷积操作组成。第一个卷积操作主要实现了通道的转换，在YOLOv10中，第二个卷积操作测试通过k=3，s=2，p=1来实现两倍下采样。

（2）C2fCIB

在讲C2fCIB之前，有必要带大家回顾一下C2f（还是简单回顾一下吧）。

C2f一般被认为是C2 faster，即更快的C2操作。如下为C2f的网络结构示意图。

C2fCIB就是将原来的Bottleneck换成了CIB模块，获得如下的结构。

CIB是一个标准的Bottleneck结构，其结构示意图如下。需要说明的是，如下是一个典型图，采取了shortcut为True，lk为False作为前提。

（3）PSA(partial self-attention)

YOLOv10引进了注意力机制，在PSA中，就集成了MHSA和FFN模块。

3. 代码解读

要想了解代码工程是如何创建模型的，还得是多关注task.py文件。

在此我们主要追随YOLOv10DetectionModel的脉络进行讲解，其他的网络构建过程基本类似。

（1）BaseModel(nn.Module)

BaseModel，从名字上就能够理解，这是一个基类。BaseModel实际上是抽象出了一些共有成员，如forward，predict，loss等，这样在后面的派生类中，只需要将更细力度的功能进行完善，就可以调用这些共有的成员进行训练。

多说一句，在面向对象编程中，不管是PyThon还是C++，都可以在基类中构造统一接口，之后通过在派生类中继承或重写，就可以实现统一的接口调用。

在BaseModel中，我们着重关注_predict_once和fuse函数。

_predict_once

_predict_once用于一次预测任务。

第128行可以看出，是将输入x依次通过self.model中的module（m）进行运算（133行）。

131行控制是否打印运行中的过程信息。

135行控制可视化。

137行控制返回嵌入特征数据。

fuse

fuse操作主要是控制一些OP的融合，以128行为例，如果m为Conv，且m中含有BN层，那么就可以按照如下进行融合，以提升计算效率。

以Conv+BN为例，我们都了解，卷积层（Conv）实际上是一个y = ax+b的过程。

训练阶段，基本的操作包含两个部分。

Conv: $x=conv.weight*x+conv.bias$

BN: $x=bn.\gamma *\frac{x_{i}-bn.mean}{\sqrt{bn.var+bn.\xi }}+bn.\beta$

其中，bn.mean为均值（对应nn.BatchNorm2d中的running_mean），bn.var为方差(对应nn.BatchNorm2d中的running_var)， $bn.\gamma$ 和 $bn.\beta$ 分别对应nn.BatchNorm2d中的weight和bias， $bn.\varepsilon$ 对应nn.BatchNorm2d中的eps。

推理阶段，可将Conv和BN进行相应的合并计算，同时需要将参数进行重新的映射。具体如下。

$x=\frac{bn.\gamma*conv.weight}{\sqrt{bn.var+bn.\varepsilon }}*x+\frac{bn.\gamma *conv.bias}{\sqrt{bn.var+bn.\epsilon }}+bn.\beta$

如上，其中 $\frac{bn.\gamma*conv.weight}{\sqrt{bn.var+bn.\varepsilon}}$ 组成了新的weight， $\frac{bn.\gamma *conv.bias}{\sqrt{bn.var+bn.\epsilon }}+bn.\beta$ 组成了新的bias。

（2）DetectionModel(BaseModel)

DetectionModel继承自BaseModel。

DetectionModel在初始化阶段主要执行三个操作，分别是定义模型，预计算参数以及初始化网络模型权重。

290行到297行主要的工作就是定义模型。其中关键的是parse_model函数，一路走来研究YOLO算法的同学们一定不陌生，后续再和大家一起回顾一下。

300行到311行则是计算stride信息，在执行模型输出解算的时候会用到。这其中实际上是构建了一组全0的数据，shape为(1, 3, 256, 256)，通过调用forward接口，进而通过网络输出来计算不同支路的stride信息。

314是进行参数的初始化。其中initialize_weights在torch_utils.py中定义。此处实际上仅仅是对BN层参数和激活层做简单的设置。其中作者还隐藏了使用nn.init.kaiming_normal_对卷积参数进行初始化。

（3）YOLOv10DetectionModel(DetectionModel)

YOLOv10DetectionModel继承自DetectionModel，仅设置了Loss计算方式，此外啥也没做，啥都用DetectionModel的。

（4）拓展讲解：parse_model

拓展重点放在模型的构建上。

前面讲到，大家对parse_model应该是不陌生的。parse_model主要是读取网络配置文件，如yolov10m.yaml，来构建生成nn.Module类型的网络。

以yolov10m.yaml为例，如下所示为文件内容。

其中，nc表示类别数，这一参数会被数据定义中的相关类别数进行覆盖。具体是在./ultralytics/models/yolov10/train.py中的YOLOv10DetectionTrainer构件时，重载get_model接口时，nc使用了self.data[“nc”]。而self.data是在./ultralytics/engine/trainer.py中的BaseTrainer中，在初始化阶段调用check_det_dataset接口获得，具体操作室读取如VOC.yaml文件获得。

scales则是模型大小的缩放。在以往的YOLO算法，如YOLOv5等，就开始提供不同尺度的网络模型，以适应不同的部署情况。因此scales就是用于将模型进行相应的缩放处理。

以第一行的Conv为例，其中第一行内容为“-1, 1, Conv, [64, 3, 2]”四个参数。从如下的parse_model局部代码截图的851行可以看到，这四个参数实际上是代表了from, number, module和args。

        from表示输入来自哪一层，-1表示相邻的上一层。

        number表示模块的重复个数，如CSP结构中Bottleneck的数量。

        module表示模块类型，此处为Conv，卷积层。

        args则是表明了当前模型的输入参数。此处为输出channel为64，使用3*3的卷积核，stride为2。