yolo各模块详解

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的类，它是nn.Module的子类。该类执行一个标准的卷积操作，可选地包括激活函数和批归一化。

该类具有以下属性和方法：

• default_act：这是一个类属性，设置为nn.SiLU()（Sigmoid-weighted Linear Unit），它是卷积操作中使用的默认激活函数。
• __init__：这是构造方法，用于初始化Conv对象。它接受以下参数：
• c1：输入通道数。
• c2：输出通道数。
• k：卷积核大小（默认为1）。
• s：步长大小（默认为1）。
• p：填充大小（默认为None）。
• g：分组数（默认为1）。
• d：膨胀因子（默认为1）。
• act：激活函数或True（默认为True）。

在这个方法中，初始化了一个Conv2d层、BatchNorm2d层和激活函数，并将它们存储为属性。

在该方法中，首先计算隐藏通道数c_，然后初始化了三个Conv层 cv1, cv2 和 cv3，其中 cv1 和 cv2 是1x1的卷积层，cv3 是具有输出通道数c2的1x1的卷积层。这些Conv层用于处理输入数据。 然后，使用nn.Sequential()初始化了一个包含n个Bottleneck模块的Sequential网络，每个Bottleneck模块的输入和输出通道数都是c_。

在该方法中，首先计算隐藏通道数c_，然后初始化了两个Conv层 cv1 和 cv2。其中 cv1 是用于处理输入张量的1x1卷积层，将输入通道数减少到c_。cv2 是用于最终输出的1x1卷积层，将输入通道数增加到c2。 接下来，使用nn.MaxPool2d初始化了一个最大池化层，池化核大小为k，步长为1，填充大小为k // 2。这个最大池化层用于执行空间金字塔池化操作。

• forward：该方法执行Conv对象的前向传播。它接受一个输入张量x，按照顺序应用卷积操作、批归一化和激活函数。然后返回输出张量。

• forward_fuse：该方法类似于forward，但不包括批归一化。它只执行卷积操作和激活函数。然后返回输出张量。

  class C3(nn.Module):
      # CSP Bottleneck with 3 convolutions
      def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
          super().__init__()
          c_ = int(c2 * e)  # hidden channels
          self.cv1 = Conv(c1, c_, 1, 1)
          self.cv2 = Conv(c1, c_, 1, 1)
          self.cv3 = Conv(2 * c_, c2, 1)  # optional act=FReLU(c2)
          self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
  
      def forward(self, x):
          return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
  

  代码定义了一个名为C3的类，它是nn.Module的子类。该类实现了一个带有3个卷积层的CSP Bottleneck结构。

  该类有以下属性和方法：

• __init__：构造方法，用于初始化C3对象。它接受以下参数：
• c1：输入通道数。
• forward：该方法执行C3对象的前向传播。它接受一个输入张量x，首先将输入张量通过cv1进行卷积，然后将输入张量通过cv2进行卷积，并将两者结果进行拼接。接着将拼接后的张量作为输入传递给Sequential网络m进行处理。最后，将处理后的结果和通过cv3进行卷积的输入张量进行拼接，并返回最终的输出张量。
  • c2：输出通道数。
  • n：重复次数（默认为1）。
  • shortcut：是否使用残差连接（默认为True）。
  • g：分组数（默认为1）。
  • e：扩展系数（默认为0.5）。

    class SPPF(nn.Module):
        # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
        def __init__(self, c1, c2, k=5):  # equivalent to SPP(k=(5, 9, 13))
            super().__init__()
            c_ = c1 // 2  # hidden channels
            self.cv1 = Conv(c1, c_, 1, 1)
            self.cv2 = Conv(c_ * 4, c2, 1, 1)
            self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
    
        def forward(self, x):
            x = self.cv1(x)
            with warnings.catch_warnings():
                warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning
                y1 = self.m(x)
                y2 = self.m(y1)
                return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
    

    码定义了一个名为SPPF的类，它是nn.Module的子类。该类实现了一个用于YOLOv5的SPPF（Spatial Pyramid Pooling - Fast）层。

    该类有以下属性和方法：

  • __init__：构造方法，用于初始化SPPF对象。它接受以下参数：
  • c1：输入通道数。
• forward：该方法执行SPPF对象的前向传播。它接受一个输入张量x，首先通过cv1对输入张量进行1x1的卷积。接下来，使用最大池化层m对cv1的输出张量进行多次池化操作，生成不同尺度的特征图。最后，将cv1的输出张量、池化特征图进行拼接，并通过cv2进行1x1的卷积，得到最终的输出张量。
  • c2：输出通道数。
  • k：池化核大小（默认为5）。

    def parse_model(d, ch):  # model_dict, input_channels(3)
        # Parse a YOLOv5 model.yaml dictionary
        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)  # redefine default activation, i.e. Conv.default_act = nn.SiLU()
            LOGGER.info(f"{colorstr('activation:')} {act}")  # print
        na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchors
        no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)
    
        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]
    
            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)
        return nn.Sequential(*layers), sorted(save)

    代码定义了一个名为parse_model的函数，用于解析YOLOv5的model.yaml字典。

    该函数接受两个参数：d是一个包含YOLOv5模型配置信息的字典，ch是输入通道数的列表。

    函数的主要功能是根据模型配置信息解析出网络的结构，并返回一个包含所有层的nn.Sequential对象以及一个已排序的保存列表。

    具体的解析过程如下：

  • 首先从模型配置信息中提取anchors，nc，gd，gw和act等值。其中，anchors是anchor框的尺寸信息，nc是类别数，gd是深度相关的参数，gw是宽度相关的参数，act是激活函数的名称。
  • 如果act存在，则将Conv类的默认激活函数default_act重新定义为act所指定的激活函数。
  • 根据anchors的信息计算出na（anchor框的数量），并计算出no（输出的特征图通道数）。
  • 初始化layers、save和c2三个列表，用于存储每一层的模块、需要保存的层的索引和最后一个层的通道数。
  • 遍历模型配置信息中的每一层，根据模块类型和参数构建相应的模块对象，并添加到layers列表中。
  • 在遍历的过程中，根据不同的模块类型，对参数进行相应的处理。例如，对于Conv、Bottleneck、SPP等模块，需要根据通道数和深度倍数进行调整；对于Detect和Segment等模块，需要将通道数的列表作为参数传入；对于Contract和Expand等模块，需要计算通道数的改变。
  • 将每个模块对象的相关信息存储到相应的属性中，并将保存的层信息添加到save列表中。
  • 最后，返回一个nn.Sequential对象，其中包含所有的模块，以及一个已排序的保存列表。