YOLOv5代码解析——模型结构篇

前言

YOLOv5🚀出到第七个版本了( •̀ ω •́ )✧，同时支持图片分类、目标检测与实例分割；我们在跑通过模型训练与推理后，可以尝试改进模型😀，或者根据任务需求来修改网络结构与损失函数等等。

本文分享一下，在模型结构方面，如何快速理解源码。

https://github.com/search?q=yolov5

一、整体代码思路（模型结构）

工程代码中，模型结构是在models目录中，其中：

• common.py 存放各个模型组件
• yolo.py 构建模型结构的主代码
• xxx.yaml 存放不同大小的模型结构配置（包括：yolov5s.yaml 、yolov5m.yaml、yolov5l.yaml、yolov5x.yaml等

它们之间的关系是：yolo.py调用common.py中的模型组件，同时解析yolov5s.yaml中的模型配置，来构建模型（backbone + head）

下面首先分享xxx.yaml的模型配置文件，再讲common.py的模型组件，最后讲yolo.py结构的主代码。

二、模型配置文件xxx.yam

这里以yolov5s.yaml为示例，它首先定义了一些模型超参数，包括类别、模型深度系数、模型宽度系数、三组初始框（anchors）。

• nc 类别：这个根据实际修改；比如COCO是80分类，nc填写80；如果是2分类，nc填写2.
• depth_multiple 模型深度系数：这是对模型的深度（层数）进行缩放，比如A模型组件中，本来是有10层卷积组成的，在构建时需要乘以这个深度系数，即：10 * 0.33 = 3.3，四舍五入后，为3层。那么实际构建时，只使用3层卷积操作。
• width_multiple 模型宽度系数：这是对模型的宽度（通道数）进行缩放，比如B模型组件中，本来的通道数为90，在构建时需要乘以这个宽度系数，即：90 * 0.50 = 45，四舍五入后，为45。那么实际构建时，通道数为45进行的。
• anchors 初始框：有大中小三组框，每一组各有三个框，参数指定了框的宽高；

# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license

# 定义一些模型超参数
nc: 80  # 类别
depth_multiple: 0.33  # 模型深度 系数
width_multiple: 0.50  # 模型宽度 系数
anchors:
  - [10,13, 16,30, 33,23]  # P3/8 用于检测小目标的三个初始框（anchors）
  - [30,61, 62,45, 59,119]  # P4/16 用于检测中目标的三个初始框（anchors）
  - [116,90, 156,198, 373,326]  # P5/32 用于检测大目标的三个初始框（anchors）

下面看看backbone 主干网络，from是指当前模块组件的输入来自那里，number是指模块重复数量，module是模块组件的名称，args是创建模型需要的参数

# YOLOv5 v6.0 backbone 主干网络
backbone:
  # [from, number, module, args] # from是指输入来自那里，number是指模块重复数量，module是模块组件的名称，args是创建模型需要的参数
  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2 第0层，相对原图做了2倍下采样
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4 第1层，相对原图做了4倍下采样
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8 第3层，相对原图做了8倍下采样
   [-1, 6, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16 第5层，相对原图做了16倍下采样
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32 第7层，相对原图做了32倍下采样
   [-1, 3, C3, [1024]],
   [-1, 1, SPPF, [1024, 5]],  # 9 第9层
  ]

from中，通常都是-1的值，它是指来自上一层；

        如果是[5, 3, C3, [128]]，from的值是5，说明该模型的输入是5-P4/16。

number，是指模块重复数量，通常都是1的值，它是指该模块只重复1遍；

        如果是[5, 3, C3, [128]]，number的值是3，说明该模块，由C3组件连续重复3遍而组成。

module，是模块组件的名称；比如C3、Conv、SPPF等（具体定义在common.py找到）

args，是创建模型需要的参数，

        比如在Conv中，对于的参数含义为：(输入通道数, 输出通道数, 卷积核大小kernel, 步长stride, 填充数padding, 分组数量groups, 扩张率dilation, 激活函数activation)

        初始值：def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True)

       比如 [-1, 1, Conv, [128, 3, 2]]，输出通道数128，卷积核大小为3，步长为2。

下面看看head 检测头，它对backbone中提取出来的特征，进一步融合，最后输出到检测头中；检测头由三个分支组成。

# YOLOv5 v6.0 head 检测头
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, C3, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small) 检测小目标 分支

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium) 检测中目标 分支

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large) 检测大目标 分支

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5) 三个检测分支
  ]

其中，解释一下 [[17, 20, 23], 1, Detect, [nc, anchors]]

[17, 20, 23] 是指Detect输入来自17层、20层、23层的特征图

1 是指 Detect模块只重复一次

Detect 是指模块的名称为Detect

[nc, anchors] 是只创建Detect模块所需的参数（类别数、三组初始框anchors）

三、模型组件common.py

common.py定义了各组模型组件，包括：

• Conv,  GhostConv,  Bottleneck,  GhostBottleneck,  SPP,  SPPF,
• DWConv,  MixConv2d,  Focus,  CrossConv, BottleneckCSP,
• C3, C3TR, C3SPP, C3Ghost,  DWConvTranspose2d, C3x
• TransformerBlock、Expand等等

通过这些组件来构建YOLO模型（backbone + head），比如看一下常用的Conv：

class Conv(nn.Module):
    # Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
    default_act = nn.SiLU()  # default activation

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()

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

    def forward_fuse(self, x):
        return self.act(self.conv(x))

它的模块组件名称为Conv，它的结构，由卷积、归一化、激化函数组成的。

模块定义所需的参数：def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True)

这和xx.yaml配置文件中的 [-1, 1, Conv, [256, 1, 1]] 参数是对应的，输出通道数256，卷积核大小为1，步长为1

四、构建模型结构主代码yolo.py

YOLOv5🚀出到第七个版本了( •̀ ω •́ )✧，同时支持图片分类、目标检测与实例分割；其中yolo.py 这个代码是模型结构的重点！！！

• DetectionModel 类，是定义目标检测的模型结构
• SegmentationModel 类，是定义实例分割的模型结构
• ClassificationModel 类，定义了图片分类的模型结构

首先重点讲一下目标检测 DetectionModel 类，后面有时间再补充实例分割、图片分类。

在yolo.py中，还能看到定义了Detect类（head 检测头）、BaseModel类（backbone 主干网络），和parse_model函数（解析xxx.yaml构建模型结构）

目标检测 DetectionModel  = （BaseModel + Detect）

4.1 parse_model函数

下面解释一下parse_model函数

它用于解析一个 YOLOv5 模型配置文件（比如：yolov5s.yaml）。函数名为 parse_model，接受两个参数 d 和 ch，其中 d 是模型配置文件的字典表示，ch 是输入图像的通道数。

思路流程：

1. 解析模型配置文件的各项参数，包括 anchors（锚框），nc（类别数目），gd（模型深度系数），gw（模型宽度系数），act（激活函数）等等。
2. 分别遍历模型配置文件中的 backbone 和 head 部分。
3. 根据配置文件中的参数，创建模型的各个层（如 Conv、Bottleneck、nn.BatchNorm2d 等）。

parse_model函数中的第一部分：

    LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10}  {'module':<40}{'arguments':<30}")
    anchors, nc, gd, gw, act = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'], d.get('activation') 
    if act:
        Conv.default_act = eval(act)  # 如果配置文件指定了激活函数，会根据配置文件的函数去加载, 比如： Conv.default_act = nn.SiLU()
        LOGGER.info(f"{colorstr('activation:')} {act}")  
    na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # anchors数量
    no = na * (nc + 5)  # 输出通道数量 = anchors * (classes + 5)

• 第一行代码：是打印日志，把yolov5s.yaml中的内容显示出来；

• 第二行代码：从模型配置文件的字典 d 中（yolov5s.yaml）获取锚点(anchors)、类别数(nc)、模型深度系数(gd)、模型宽度系数(gw)、激活函数(act)等参数，并将其赋值给相应的变量 anchors、nc、gd、gw、act。

• 第三行代码：如果配置文件中指定了激活函数(act)，则将其通过 eval() 函数重新定义为 Conv.default_act，即重定义默认的激活函数为 nn.SiLU()，并使用日志记录 LOGGER.info() 输出激活函数的信息。

• 第六行代码：计算出锚点的数量 na。如果锚点是一个列表，则取第一个锚点列表的长度的一半作为锚点的数量；如果锚点是一个数值，则直接使用该数值作为锚点的数量。

• 第七行代码：根据锚点数量、类别数等参数计算出模型输出层的通道数量 no，即 anchors * (classes + 5)，其中 anchors 是锚点的数量，classes 是类别数，5 是包括目标置信度、边界框坐标等信息的预测数目。

parse_model函数中的第二部分：

    layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
    for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  # from, number, module, args
        m = eval(m) if isinstance(m, str) else m  # eval strings
        for j, a in enumerate(args):
            with contextlib.suppress(NameError):
                args[j] = eval(a) if isinstance(a, str) else a  # eval strings

这段代码主要是对一个包含多层网络模块的列表进行循环遍历，分别遍历backbone 和 head。

• 定义了三个变量 layers、save 和 c2，并分别初始化为空列表 []，空列表 [] 和列表 ch 的最后一个元素 ch[-1]。

• 使用 for 循环遍历 d['backbone'] + d['head'] 列表中的元素，每个元素都表示一个网络模块。其中 f 表示 "from"（该模块输入来自哪里），n 表示 "number"（模块的重复次数），m 表示 "module"（模块的名称），args 表示 "arguments"（创建模块所需的参数）。

• 对于每个网络模块，将 m （模块的名称）通过 eval() 函数进行评估操作，如果 m 是字符串类型，则将其转换为相应的 Python 对象。这里使用 eval() 函数来将字符串类型的模块名转换为实际的模块对象，从而可以在后续的代码中使用。

• 对于 args 列表中的每个元素 a，同样使用 eval() 函数进行评估操作，将字符串类型的参数值转换为相应的 Python 对象。这里使用 with contextlib.suppress(NameError) 来捕获可能出现的 NameError 异常，以防止字符串参数无法成功评估的情况。

parse_model函数中的第三部分：

    layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
    for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  # from, number, module, args
        m = eval(m) if isinstance(m, str) else m  # eval strings
        for j, a in enumerate(args):
            with contextlib.suppress(NameError):
                args[j] = eval(a) if isinstance(a, str) else a  # eval strings

#主要讲下面的代码
        n = n_ = max(round(n * gd), 1) if n > 1 else n  # depth gain
        if m in {
                Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
                BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x}:
            c1, c2 = ch[f], args[0]
            if c2 != no:  # if not output
                c2 = make_divisible(c2 * gw, 8)

            args = [c1, c2, *args[1:]]
            if m in {BottleneckCSP, C3, C3TR, C3Ghost, C3x}:
                args.insert(2, n)  # number of repeats
                n = 1
        elif m is nn.BatchNorm2d:
            args = [ch[f]]
        elif m is Concat:
            c2 = sum(ch[x] for x in f)
        # TODO: channel, gw, gd
        elif m in {Detect, Segment}:
            args.append([ch[x] for x in f])
            if isinstance(args[1], int):  # number of anchors
                args[1] = [list(range(args[1] * 2))] * len(f)
            if m is Segment:
                args[3] = make_divisible(args[3] * gw, 8)
        elif m is Contract:
            c2 = ch[f] * args[0] ** 2
        elif m is Expand:
            c2 = ch[f] // args[0] ** 2
        else:
            c2 = ch[f]

这段代码是对模型的网络层进行解析和处理的部分。

1. 首先对输入的深度参数 n 进行计算得到新的深度 n_，计算方式为将 n* gd（depth_multiple）并进行四舍五入，但至少为 1。
2. 然后根据不同的模块类型 m 进行不同的处理：
   • 如果 m 是 Conv、GhostConv、Bottleneck、GhostBottleneck、SPP、SPPF、DWConv、MixConv2d、Focus、CrossConv、BottleneckCSP、C3、C3TR、C3SPP、C3Ghost、nn.ConvTranspose2d、DWConvTranspose2d、C3x 中的一种，需要更新参数 args，并进行一些处理。
   • 如果 m 是 nn.BatchNorm2d，只需要更新参数 args。
   • 如果 m 是 Concat，需要计算并更新参数 c2。
   • 如果 m 是 Detect 或 Segment，需要更新参数 args，并进行一些处理。
   • 如果 m 是 Contract，需要根据 args[0] 计算并更新参数 c2。
   • 如果 m 是 Expand，需要根据 args[0] 计算并更新参数 c2。
   • 对于其他的模块类型，直接更新参数 c2。

比如，A模型组件中，本来是有10层卷积组成的（n = 10），在构建时需要乘以这个深度系数0.33（gd = 0.33），即：10 * 0.33 = 3.3，四舍五入后，为3层。那么实际构建时，只使用3层卷积操作。

parse_model函数中的第四部分：

        m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
        t = str(m)[8:-2].replace('__main__.', '')  # module type
        np = sum(x.numel() for x in m_.parameters())  # number params
        m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
        LOGGER.info(f'{i:>3}{str(f):>18}{n_:>3}{np:10.0f}  {t:<40}{str(args):<30}')  # print
        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
        layers.append(m_)
        if i == 0:
            ch = []
        ch.append(c2)

这段代码主要涉及对模型的不同类型的处理，并将处理后的模型添加到layers列表中，并更新一些变量

1. m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)：根据n的值判断是否需要将m模型包装在nn.Sequential中，然后将m模型与args参数传递给m的构造函数，生成一个新的模型m_。

2. t = str(m)[8:-2].replace('__main__.', '')：将模型m的类型转换成字符串，并进行字符串处理，提取模型类型信息，将模型类型保存在变量t中。

3. np = sum(x.numel() for x in m_.parameters())：计算模型m_的参数总数，将结果保存在变量np中。

4. m_.i, m_.f, m_.type, m_.np = i, f, t, np：将模型的索引i、'from'索引f、类型t和参数数量np保存在模型m_的属性中。

5. LOGGER.info(f'{i:>3}{str(f):>18}{n_:>3}{np:10.0f} {t:<40}{str(args):<30}')：打印模型的信息，包括索引i、'from'索引f、重复次数n_、参数数量np、类型t和参数args。

6. save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)：将f中不为-1的值加入到savelist中，其中如果f是整数，则将其作为列表处理。

7. layers.append(m_)：将处理后的模型m_添加到layers列表中。

8. if i == 0: ch = []：如果当前模型的索引i为0，则将ch列表清空。

9. ch.append(c2)：将变量c2添加到ch列表中。

4.2 BaseModel 类

BaseModel 类是用来构建backbone的，定义了一些方法：

• 用于前向推理计算（forward、_forward_once），用于单尺度推理和训练时的前向传播计算。
• 单层性能分析（_profile_one_layer），用于分析每个层的性能。
• 模型融合（fuse），用于融合模型中的 Conv2d() 和 BatchNorm2d() 层，以减少计算量。
• 模型信息打印（info），用于打印模型的信息，包括层的类型、输入输出维度、运算时间、浮点运算次数等。
• 张量转换（_apply），用于对模型中的张量应用转换，如 to()、cpu()、cuda()、half() 等操作。

在前向推理计算中，forward函数会调用_forward_once和_profile_one_layer进行前向推理的模型构建。

• 在前向计算方法 _forward_once 中，模型中的每个层被遍历并执行运算。如果一个层的输入不是来自于前一层，则输入是来自于模型中之前的某一层。
• 其中使用self.model列表存储了模型中的所有层，_forward_once函数则遍历所有层，按照顺序逐层计算并将每层的输出保存到列表y中。
• 如果设置了 profile，则会调用 _profile_one_layer 方法来对当前层进行性能分析，记录时间和浮点运算次数。在每一层的输出被保存后，如果设置了可视化，则会将当前层的特征图保存到指定路径下。

class BaseModel(nn.Module):
    # YOLOv5 base model
    def forward(self, x, profile=False, visualize=False):
        return self._forward_once(x, profile, visualize)  # single-scale inference, train

    def _forward_once(self, x, profile=False, visualize=False):
        y, dt = [], []  # outputs
        for m in self.model:
            if m.f != -1:  # if not from previous layer
                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
            if profile:
                self._profile_one_layer(m, x, dt)
            x = m(x)  # run
            y.append(x if m.i in self.save else None)  # save output
            if visualize:
                feature_visualization(x, m.type, m.i, save_dir=visualize)
        return x

单层性能分析（_profile_one_layer），用于分析每个层的性能。该函数使用thop库计算每层的FLOPs，使用time_sync函数计算每层的时间消耗，并输出每层的时间、FLOPs和参数信息。

    def _profile_one_layer(self, m, x, dt):
        c = m == self.model[-1]  # is final layer, copy input as inplace fix
        o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1E9 * 2 if thop else 0  # FLOPs
        t = time_sync()
        for _ in range(10):
            m(x.copy() if c else x)
        dt.append((time_sync() - t) * 100)
        if m == self.model[0]:
            LOGGER.info(f"{'time (ms)':>10s} {'GFLOPs':>10s} {'params':>10s}  module")
        LOGGER.info(f'{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f}  {m.type}')
        if c:
            LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s}  Total")

下面是模型融合（fuse）、模型信息打印（info）、张量转换（_apply）

    # 用于融合模型中的 Conv2d() 和 BatchNorm2d() 层，以减少计算量
    def fuse(self):  
        LOGGER.info('Fusing layers... ')
        for m in self.model.modules():
            if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
                m.conv = fuse_conv_and_bn(m.conv, m.bn)  
                delattr(m, 'bn')  
                m.forward = m.forward_fuse  
        self.info()
        return self

    # 用于打印模型的信息，包括层的类型、输入输出维度、运算时间、浮点运算次数等。
    def info(self, verbose=False, img_size=640):  
        model_info(self, verbose, img_size)

    # 用于对模型中的张量应用转换，如 to()、cpu()、cuda()、half() 等操作
    def _apply(self, fn):
        self = super()._apply(fn)
        m = self.model[-1]  # Detect()
        if isinstance(m, (Detect, Segment)):
            m.stride = fn(m.stride)
            m.grid = list(map(fn, m.grid))
            if isinstance(m.anchor_grid, list):
                m.anchor_grid = list(map(fn, m.anchor_grid))
        

4.3 Detect类

用于构建YOLOv5目标检测模型中的检测头部，负责生成检测结果。其中类中有以下三个函数：

• __init__(self, nc=80, anchors=(), ch=(), inplace=True): 初始化函数，用于构建 Detect 类的对象。nc 表示类别数量，anchors 表示锚框的坐标，ch 表示输入特征图的通道数，inplace 表示是否使用原地（inplace）操作。该方法会初始化模型的各个属性，并构建模型的卷积层。
• forward(self, x): 前向传播函数，用于生成检测结果。输入参数 x 是一个包含多个特征图的列表。在前向传播过程中，模型会对输入的每个特征图进行卷积操作，然后根据检测结果的格式，生成相应的输出。如果处于训练模式，返回卷积层的输出；如果处于推理模式，返回生成的检测结果。
• _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, '1.10.0')): 用于生成网格坐标和锚框在特征图上的坐标。nx 和 ny 分别表示特征图的宽度和高度，i 表示当前特征图的索引，torch_1_10 是一个用于检查 torch 版本的函数，用于选择不同的方式生成网格坐标，以保持兼容性。该方法会返回网格坐标和锚框在特征图上的坐标。

看一下__init__中的变量含义：

class Detect(nn.Module):
    stride = None  # 检测层的步长
    dynamic = False  # 表示是否强制进行网格重构
    export = False  # 表示是否处于导出模式

    def __init__(self, nc=80, anchors=(), ch=(), inplace=True):  # detection layer
        super().__init__()
        self.nc = nc  # 类别数
        self.no = nc + 5  # 每个锚框的输出通道数，等于类别数加上5（x、y、w、h、confidence）
        self.nl = len(anchors)  # 检测层的数量
        self.na = len(anchors[0]) // 2  # 每个检测层的锚框数量
        self.grid = [torch.empty(0) for _ in range(self.nl)]  # 用于存储生成的网格，初始值为空tensor
        self.anchor_grid = [torch.empty(0) for _ in range(self.nl)]  # 用于存储生成的锚框，初始值为空tensor
        self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2))  # shape(nl,na,2)
        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # 用于输出检测结果的卷积操作
        self.inplace = inplace  # 是否使用inplace操作

Detect类的前向传播函数forward接收一个输入x，其中x是一个列表，包含了多个尺度的特征图（3个分支）。在前向传播过程中，对每个特征图进行处理（检测头部分）：

    def forward(self, x):
        z = []  # inference output
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

            if not self.training:  # inference
                if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)

                if isinstance(self, Segment):  # (boxes + masks)
                    xy, wh, conf, mask = x[i].split((2, 2, self.nc + 1, self.no - self.nc - 5), 4)
                    xy = (xy.sigmoid() * 2 + self.grid[i]) * self.stride[i]  # xy
                    wh = (wh.sigmoid() * 2) ** 2 * self.anchor_grid[i]  # wh
                    y = torch.cat((xy, wh, conf.sigmoid(), mask), 4)
                else:  # Detect (boxes only)
                    xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4)
                    xy = (xy * 2 + self.grid[i]) * self.stride[i]  # xy
                    wh = (wh * 2) ** 2 * self.anchor_grid[i]  # wh
                    y = torch.cat((xy, wh, conf), 4)
                z.append(y.view(bs, self.na * nx * ny, self.no))

        return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)

代码解析：

1. 初始化一个空列表 z，用于保存推断（inference）输出。
2. 进行一个循环，循环次数为 self.nl，self.nl 是模型中的层数。
3. 对每一层进行前向传播计算：
   • 调用 self.m[i](x[i]) 对输入 x[i] 进行卷积计算，其中 self.m[i] 是第 i 层的卷积层。
   • 将 x[i] 的形状从 (bs,255,20,20) 转换为 (bs,3,20,20,85)，其中 bs 是 batch size，255 是特定的通道数，20 是特定的高度和宽度，85 是特定的预测目标数量和属性数。
   • 使用 view 函数将 x[i] 进行形状变换，变换为 (bs, self.na, self.no, ny, nx)，其中 self.na 是 anchor boxes 的数量，self.no 是预测的目标属性数量，ny 和 nx 是特定的高度和宽度。
   • 使用 permute 函数对维度进行重排列，变换为 (bs, 1, ny, nx, self.no)，其中第二个维度 1 对应于 anchor boxes 的数量 self.na。
   • 使用 contiguous 函数使数据在内存中连续存储，以便进行后续计算。
   • 根据是否处于训练模式（self.training），进行不同的推断处理：
     • 如果处于推断模式（self.training 为 False）：
       • 检查是否需要动态调整 anchor boxes 的网格大小或者 anchor boxes 是否需要重新生成，如果需要，则调用 _make_grid 函数进行生成。
       • 根据模型是否属于 Segment 类进行不同的处理：
         • 如果是 Segment 类，则将 x[i] 进行分割，分别得到 xy、wh、conf 和 mask，分别表示预测的目标的中心坐标、宽高、置信度和掩码（分割结果）。
         • 如果不是 Segment 类，那些是Detect，则将 x[i] 进行目标检测，分别得到 xy、wh 和 conf，分别表示预测的目标的中心坐标、宽高和置信度。
         • 对 xy 进行 sigmoid 函数计算，然后乘以 2 并加上网格坐标，再乘以 self.stride[i] 进行缩放，得到预测的目标的中心坐标 xy。
         • 对 wh 进行 sigmoid 函数计算，然后将结果平方并乘以 self.anchor_grid[i] 进行缩放，得到预测的目标的宽高 wh。
         • 使用 torch.cat 函数将 xy、wh、conf.sigmoid() 和 mask 拼接在一起，形成预测的目标信息 y。
         • 将 y 进行形状变换，变换为 (bs, self.na * nx * ny, self.no)，其中 bs 是 batch size，self.na 是 anchor boxes  

解析代码1： Segment:  (boxes + masks)

xy, wh, conf, mask = x[i].split((2, 2, self.nc + 1, self.no - self.nc - 5), 4)

这段代码使用了 PyTorch 中的 split 函数，用于将张量 x[i] 沿着指定的维度（在这里是第4维，通道维度）进行切分，并将切分后的子张量赋值给变量 xy, wh, conf, 和 mask。

x[i] 是一个形状为 (bs, na, ny, nx, no) 的张量，其中：

• bs 表示 batch size，即批次大小；
• na 表示每个位置的 anchor 数量；（每个网格默认3个）
• ny 表示特征图的高度；
• nx 表示特征图的宽度；
• no 表示每个 anchor 预测的输出通道维度数量。

split 函数接受一个元组作为参数，其中包含了要切分的张量 x[i] 在第4维上（通道维度）的切分位置，以及切分后的子张量的数量。在这里，切分位置是 (2, 2, self.nc + 1, self.no - self.nc - 5)，表示从第4维的索引0开始，切分长度分别为2、2、self.nc + 1、self.no - self.nc - 5，总共切分为4个子张量。

切分后的子张量分别赋值给 xy, wh, conf, 和 mask，其中：

• xy 是包含了预测的边界框在特征图上的中心坐标的张量；
• wh 是包含了预测的边界框在特征图上的宽度和高度的张量；
• conf 是包含了预测的边界框的置信度的张量；
• mask 是包含了预测的边界框的分割掩码信息的张量。

这段代码的作用是从输入张量 x[i] 中将预测的边界框的中心坐标、宽度和高度、置信度以及掩码信息分别提取出来，并赋值给相应的变量，以便后续处理和使用。

解析代码2：

为什么是self.no - self.nc - 5？

self.no 表示每个 anchor 预测的输出的通道维度数量，而 self.nc 则表示每个 anchor 预测的类别数量。

在这段代码中，使用了 self.no - self.nc - 5 作为切分的长度，是因为在 YOLO 模型中，每个 anchor 预测的输出包括了边界框的位置信息（中心坐标、宽度和高度）、是否包含物体的置信度以及类别概率。其中，边界框位置信息需要4个维度来表示（2个维度表示中心坐标，2个维度表示宽度和高度），置信度需要1个维度来表示，而类别概率根据类别数量 self.nc 而定。

所以，self.no - self.nc - 5 表示在剩余的维度中，预测的边界框掩码信息的长度，用于将输入张量 x[i] 切分成 xy, wh, conf, 和 mask 这四个子张量。这个长度的计算是基于 YOLO 模型的设计和输出维度的特点，可能在不同的模型中会有不同的值。

解析代码3：Detect (boxes only)

xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4)

这段代码使用了 PyTorch 中的 split 函数，用于将张量 x[i] 沿着指定的维度（在这里是第4维，通道维度）进行切分，并将切分后的子张量赋值给变量 xy, wh, conf。

x[i] 是一个形状为 (bs, na, ny, nx, no) 的张量，其中：

• bs 表示 batch size，即批次大小；
• na 表示每个位置的 anchor 数量；（每个网格默认3个）
• ny 表示特征图的高度；
• nx 表示特征图的宽度；
• no 表示每个 anchor 预测的输出通道维度数量。

 xy = (xy * 2 + self.grid[i]) * self.stride[i]

这段代码是对目标框的中心坐标进行计算和转换的操作。

• xy 是通过 x[i] 中的前两个通道进行切片得到的目标框的中心坐标预测值，其形状为 (bs, ny, nx, 2)，其中 bs 是 batch size，ny 和 nx 分别是输入特征图的高度和宽度，2 表示中心坐标的 x 和 y 分量。
• self.grid[i] 是预先计算的网格坐标偏移值，其形状与 xy 相同，用于将目标框的中心坐标从特征图空间映射到输入图像空间。
• self.stride[i] 是预先定义的特征图相对于输入图像的步长，用于将目标框的中心坐标进行缩放。

这段代码的计算过程为：

1. 将 xy 乘以 2，然后加上 self.grid[i]，实现从特征图空间到输入图像空间的映射。
2. 将结果乘以 self.stride[i]，实现对目标框中心坐标的缩放。

最终，xy 存储了转换后的目标框中心坐标值。

wh = (wh * 2) ** 2 * self.anchor_grid[i] 

这段代码计算了预测框的宽高信息。

• 首先，wh乘以2，再将其平方，最后乘以self.anchor_grid[i]。
• self.anchor_grid[i]是一个锚框（anchor box）的尺寸，用于调整预测框的宽高。
• 在目标检测算法中，锚框是一些预定义的框，用于在不同尺度和长宽比下对图像进行采样和预测。通常情况下，锚框的尺寸和比例是根据数据集和任务需求进行设置的。

这段代码的目的是对预测框的宽高信息进行调整和缩放，以便与图像实际尺寸相匹配。最终的wh张量将包含经过缩放后的预测框的宽高信息，用于后续的目标检测任务。

y = torch.cat((xy, wh, conf), 4)

这段代码使用torch.cat()函数将xy、wh和conf三个张量在第4维度（即通道维度）上进行拼接，生成一个新的张量y。

这三个张量在这段代码中按顺序进行拼接，生成一个新的张量y，其中包含了预测框的信息，包括预测框的中心点坐标、宽高和置信度。

在forward函数中会调用_make_grid函数，它用于生成 anchor boxes 在输入图像上的网格。

    def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, '1.10.0')):
        d = self.anchors[i].device
        t = self.anchors[i].dtype
        shape = 1, self.na, ny, nx, 2  # grid shape
        y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
        yv, xv = torch.meshgrid(y, x, indexing='ij') if torch_1_10 else torch.meshgrid(y, x)  # torch>=0.7 compatibility
        grid = torch.stack((xv, yv), 2).expand(shape) - 0.5  # add grid offset, i.e. y = 2.0 * x - 0.5
        anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
        return grid, anchor_grid

函数的参数含义：

• nx 和 ny 分别表示在 x 和 y 方向上的网格数量，默认值为 20。
• i 表示当前 anchor boxes 的索引，默认值为 0。
• torch_1_10 是一个布尔值，用于检查 torch 版本是否大于等于 1.10.0，这是通过调用 check_version() 函数来实现的。

函数思路流程：

• 首先，函数根据 self.anchors[i] 的设备和数据类型创建了 y 和 x 张量，分别表示 y 和 x 方向上的网格索引。
• 接着，通过调用 torch.meshgrid() 函数生成了网格张量 yv 和 xv，其中 ij 索引方式在 torch 版本大于等于 1.10.0 时生效，否则使用默认的索引方式。这里生成了一个二维的网格，其中 yv 表示 y 方向上的网格索引，xv 表示 x 方向上的网格索引。
• 然后，通过调用 torch.stack() 函数将 xv 和 yv 沿着第三个维度（索引从 0 开始）堆叠在一起，得到一个形状为 (ny, nx, 2) 的网格张量。
• 接着，网格张量 grid 被扩展为与 anchor boxes 数量和维度相同的形状 (1, self.na, ny, nx, 2)，并且每个网格点都减去了 0.5 的偏移量，即 grid = torch.stack((xv, yv), 2).expand(shape) - 0.5。
• 最后，anchor boxes 张量 self.anchors[i] 被乘以对应的步长 self.stride[i]，并且形状被调整为 (1, self.na, 1, 1, 2)，然后扩展为与网格张量相同的形状 (1, self.na, ny, nx, 2)，得到了 anchor boxes 在输入图像上的坐标。
• 函数返回了两个张量，分别是 grid 和 anchor_grid，它们分别表示 anchor boxes 在输入图像上的网格坐标和相对于输入图像的坐标。

其中，x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

思路流程：

1. 使用.view(bs, self.na, self.no, ny, nx)对x[i]进行形状变换，将其变为5维张量，其中bs是batch size，self.na是每个grid cell中anchor的数量，self.no是每个anchor预测的输出通道数，ny和nx分别是grid的高度和宽度。

2. 使用.permute(0, 1, 3, 4, 2)对5维张量进行维度置换，将最后一个维度（2）移到第五个维度位置，从而变为(bs, self.na, ny, nx, self.no)的形状。

3. 使用.contiguous()将维度连续化，以确保后续操作的正确性。

最终，x[i]的形状被变换为(bs, self.na, ny, nx, self.no)，在这个形状下，可以方便地对anchor的预测信息进行处理和解析。

4.4 DetectionModel类

还没写完，待完善更新~