YOLOv5第Y4周：common.py文件解读

本文链接：https://blog.csdn.net/geo436872/article/details/134990833

🍨 本文为🔗365天深度学习训练营中的学习记录博客
🍖 原作者：K同学啊 | 接辅导、项目定制
🚀 文章来源：K同学的学习圈子

一、基本组件

1.autopad

def autopad(k, p=None, d=1):  # kernel, padding, dilation
    # Pad to 'same' shape outputs
    if d > 1:
        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p

autopad 函数在 YOLOv5 等计算机视觉系统中通常用于自动计算卷积层所需的填充（padding），以保持特征图（feature map）的空间维度或者使其符合某种设计意图，如保留边界信息或者使得卷积操作的输出大小等同于输入大小（即所谓的 ‘same’ padding）。

这个函数的作用是，给定一个卷积核尺寸 k、填充 p 和扩张率 d，它会返回一个计算出来的填充大小 p。这样做的目的在于简化网络设计过程，并使得特征图的大小在卷积运算后保持不变，或者是依照设计者的其他特定需求。

函数的具体逻辑如下：

如果卷积层使用了扩张率（dilation）d 大于 1，实际的卷积核尺寸将扩大。扩张卷积通过在卷积核的每个元素之间添加空间（比如空格或者零填充）来增大卷积核的感受野而不增加参数数量。
```
k = d * (k - 1) + 1
```
如果 k 是整数（即卷积核为正方形），直接计算；如果 k 是列表（即卷积核可能不是正方形，比如矩形），则对每个维度分别计算。
如果没有指定填充 p，函数将根据卷积核尺寸 k 来计算 ‘same’ padding 的大小，保证输出的特征图大小不变。对于整数尺寸，将卷积核大小 k 整除 2 即可得到 ‘same’ padding 的大小；如果 k 是列表，则对每个维度分别进行操作。
```
p = k // 2
```
函数返回计算出的填充大小 p。

在 YOLOv5 模型的构建中，正确的 padding 确保特征图尺寸的一致性有助于维护不同层输出之间的对齐，这一点对于模型检测性能至关重要，因为 YOLOv5 依赖不同尺度的特征图来检测不同大小的物体。

2.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))

在 YOLOv5 网络设计中，Conv 类是一个自定义层，它定义了一个标准的卷积操作，可能会包含一些额外的操作如批量标准化（Batch Normalization）和激活函数。以下是 Conv 类定义中各部分的详细解释：

class Conv(nn.Module)：nn.Module 是 PyTorch 中所有神经网络层的基类。Conv 继承自 nn.Module，表示它是一个自定义的卷积层。
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True)：这是类的构造函数，用于初始化层的状态。参数含义如下：
- c1：输入通道数。
- c2：输出通道数。
- k：卷积核大小，默认为 1。
- s：步长（stride），默认为 1。
- p：填充（padding），如果为 None，则自动计算以实现 ‘same’ padding。
- g：组数（groups），用于分组卷积，默认为 1，表示不分组。
- d：扩张率（dilation），默认为 1。
- act：激活函数，如果为 True，则使用默认激活函数 SiLU；如果是 nn.Module 的实例，则使用该激活函数；如果不是，则使用恒等映射（即没有激活函数）。
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)：创建一个二维卷积层 Conv2d，其参数已经在构造函数中定义。autopad 函数用于计算卷积层所需的填充大小以保持特征图尺寸。bias=False 表示该卷积层不使用偏置项。
self.bn = nn.BatchNorm2d(c2)：创建一个批量标准化层，用于标准化特征图的输出以提高训练的稳定性和网络的泛化性能。
self.act = ...：初始化激活函数，基于 act 参数确定使用哪种激活函数。
def forward(self, x)：定义前向传播函数。输入张量 x 通过卷积层，然后是批量标准化层，最后应用激活函数，最终返回处理后的输出。
def forward_fuse(self, x)：定义了另一种前向传播函数，可能用于简化模型（将批量标准化参数融合到卷积权重中）。在这个函数中，输入直接通过卷积层，并应用激活函数，省略了批量标准化层的使用。

3.Focus

class Focus(nn.Module):
    # Focus wh information into c-space
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super().__init__()
        self.conv = Conv(c1 * 4, c2, k, s, p, g, act=act)
        # self.contract = Contract(gain=2)
 
    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))
        # return self.conv(self.contract(x))

Focus 类在 YOLOv5 模型中起着空间信息收缩（focusing）的作用。它被设计来缩减输入特征图的宽度和高度（即空间分辨率），同时将通道信息扩增。此操作有助于在不损失重要信息的情况下减少后续层需要处理的数据。让我们细致地看下 Focus 类的构造和作用：

__init__ 方法初始化 Focus 类的一个实例。其参数与 Conv 类似，并构建了一个 Conv 对象。这个卷积是为了在特征图被“focusing”之后进行进一步处理。
- c1 * 4 因为 Focus 层会将输入张量的通道数增加四倍（解释在 forward 方法中），因此传递给 Conv 类的输入通道数是 c1 * 4。
- c2 是 Conv 层的输出通道数。
forward 方法中，Focus 类将输入张量 x 分为四个相同大小的部分：
- x[..., ::2, ::2] 选取了从每个角开始，每隔一个像素点的子集。
- x[..., 1::2, ::2] 选取了第一行（从第二个像素开始）的每隔一个像素点，并对所有偶数行应用此操作。
- x[..., ::2, 1::2] 选取偶数列的像素点，并应用于所有的行。
- x[..., 1::2, 1::2] 从第二行第二列的像素开始，保留行和列的每隔一个像素点。

将这四个部分在通道维度（第二维度）上拼接，使得最终输出的通道数是输入的四倍。这个“focusing”过程降低了空间分辨率（宽度和高度都变成了原来的一半），从而聚焦于更重要的特征。之后 self.conv （即初始化中已经定义好的 Conv 层）会用来进一步处理这个四倍通道的合并特征图。

4.Bottleneck

class Bottleneck(nn.Module):
    # Standard bottleneck
    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
        super().__init__()
        c_ = int(c2 * e)  # hidden channels
        # 1*1卷积层
        self.cv1 = Conv(c1, c_, 1, 1)
        # 3*3卷积层
        self.cv2 = Conv(c_, c2, 3, 1, g=g)
        # shortcut=True 并且 c1==c2 才能做shortcut, 将输入和输出相加之后再输出
        self.add = shortcut and c1 == c2
 
    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))

Bottleneck 类在 YOLOv5 架构中是一个基本的构建块，用于创建深度神经网络中的残差学习路径。残差块可以帮助网络训练更深层次的架构而不会出现梯度消失的问题。现在，让我们看看 Bottleneck 类如何实现。

__init__ 方法初始化一个 Bottleneck 类的实例。这个方法接受以下参数：
- c1：输入通道数。
- c2：输出通道数。
- shortcut：一个布尔值，指示是否使用快捷连接（即残差连接）。
- g：分组卷积中的组数，可以用于创建组卷积，这有助于模型学习更复杂的特征表示。
- e：扩展因子，用于调整隐藏层的通道数。通常，扩展因子小于 1，意味着隐藏层的通道数小于输出通道数，这样可以减少参数数量。
cv1 是一个使用 1x1 卷积核的 Conv 类实例，用于减少通道数（也称为通道压缩），从而限制模型复杂度，具有类似瓶颈的作用。
cv2 是一个使用 3x3 卷积核的 Conv 类实例，用于决定最终的输出通道数。这里使用 g 参数，可以让该卷积层以组卷积的方式运行。
self.add 确定在前向传播中是否将输入加到卷积层的输出上。这是典型残差连接的特征，允许网络直接传递输入，与经过转换的输出相加，有利于梯度的反向传播，让网络更容易学习身份映射。
forward 方法定义了 Bottleneck 的前向传播规则。如果 self.add 为 True 并且输入输出通道数相同，则输入 x 会通过两个卷积层 (cv1 和 cv2) 的处理，然后将处理后的结果与原始输入 x 相加作为输出。这就是残差连接的核心。如果不能进行快捷连接（self.add 为 False），则输出仅为两个卷积处理后的结果。

言而总之，Bottleneck 类通过使用 1x1 卷积来减少通道数，再通过 3x3 卷积来使用分组卷积处理这些通道，适用于构建深度网络层并提高网络效能。残差连接允许网络更加容易地进行训练，尤其是在训练深层网络时，它有助于信息传播和避免梯度消失问题。

5.BottleneckCSP

class BottleneckCSP(nn.Module):
    # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
    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 = nn.Conv2d(c1, c_, 1, 1, bias=False)
        self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
        self.cv4 = Conv(2 * c_, c2, 1, 1)
        self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
        self.act = nn.SiLU()
        self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
 
    def forward(self, x):
        y1 = self.cv3(self.m(self.cv1(x)))
        y2 = self.cv2(x)
        return self.cv4(self.act(self.bn(torch.cat((y1, y2), 1))))

BottleneckCSP 类是用在 YOLOv5 模型中的另一种残差结构，其灵感来源于 Cross Stage Partial Networks (CSPNet)。这种设计旨在将特征图分为两个部分，并分别进行处理，最后在深度上合并，以降低计算成本并提高网络性能。接下来，我们详细解析 BottleneckCSP 类的组成和作用：

__init__ 方法初始化这个类的实例。其接受参数如下：
- c1 是输入通道数。
- c2 是输出通道数。
- n 是在BottleneckCSP中Bottleneck块的数量。
- shortcut 表明是否在Bottleneck块中使用残差连接。
- g 是分组卷积中的组数。
- e 是扩张率，用于计算隐藏通道(c_)。
cv1 是一个利用 1x1 卷积核减少特征通道数的卷积层。
cv2 是一个不带偏置的 1x1 卷积层，它直接作用在输入 x 上，以创建的一部分特征图将会与 cv3 的输出相结合。
cv3 也是一个不带偏置的 1x1 卷积层，它作用于 m 模块的输出。m 是一个 Bottleneck 层序列，用于深入特征提取。
cv4 是一个 1x1 卷积层，用于将由 cv2 和 cv3 结合的特征通道数从 2 * c_ 降到输出通道数 c2。
bn 是批量归一化层，应用于 cv2 和 cv3 的输出合并后的结果。
act 是激活层，此处使用 SiLU (Sigmoid Linear Unit)，也称为 Swish 激活函数。
m 是一个序列（nn.Sequential），由 n 个 Bottleneck 模块组成，该模块在输入特征图 self.cv1(x) 上重复应用提取特征。

forward 方法描述了信息如何流经 BottleneckCSP 模块：

输入 x 首先通过 cv1 进行处理，得到的输出传递给由 n 个 Bottleneck 块组成的 m 序列，最终通过 cv3 得到 y1。
输入 x 同时经过 cv2 直接处理得到 y2。
接着 y1 和 y2 在通道维度上进行拼接，并通过批量归一化和 Swish 激活。
最后通过 cv4 进行 1x1 卷积，以降低特征图的通道数至所需的输出 c2。

这个 BottleneckCSP 结构可以提高特征的再利用效率，同时避免了特征的重复计算，使模型在效率和性能之间取得了良好的平衡。这是实现更深、但计算效率高的网络架构的关键元素之一。

6.C3

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
        
        # 3个1*1卷积层的堆叠，比BottleneckCSP少一个
        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类在 YOLOv5 中是BottleneckCSP的一个变种，这个模块结合了三个卷积层和若干个Bottleneck 模块。C3模块与BottleneckCSP的不同之处在于它直接使用了三个卷积层，而不是在BottleneckCSP中的四个卷积层，并且在合并特征之前没有使用批量归一化和激活函数。下面是C3` 模块的详细作用和结构：

__init__ 方法初始化这个类实例，并接受以下参数：
- c1 是输入通道数。
- c2 是输出通道数。
- n 是重复 Bottleneck 块的数量。
- shortcut 表示是否在 Bottleneck 块中使用残差连接。
- g 是分组卷积的组数。
- e 是扩张率，它用于决定隐藏通道数 c_。
cv1 是输入特征图 x 的 1x1 卷积，用于扩展/减少特征通道数。
cv2 也是对原始输入 x 的 1x1 卷积，但它的输出将会与 m 处理后的特征图合并。
cv3 是最后的 1x1 卷积，用于将合并后的特征图（来自 cv1/m 和 cv2）的通道数从 2 * c_ 减少到输出通道数 c2。
m 是一个由 n 个 Bottleneck 层组成的 nn.Sequential 序列，用于进行特征提取。

在 forward 方法中：

输入 x 被传递至 cv1 和 m 模块进行一系列的 Bottleneck 操作，得到特征 y1。
同时，x 也通过 cv2 转换，得到另一组特征 y2。
最后，y1 和 y2 在通道维度合并，接着通过 cv3 进行卷积操作产生最终输出。

C3 模块的设计使得信息可以通过不同路径流动，这样做可以加强特征的混合，改善学习效果和网络的表现力，同时因为减少了一个卷积层，也在一定程度上减少了模型参数和计算量。这样的设计有助于 YOLOv5 在提升目标检测性能的同时保持较高效的计算速度。

7.SPP

class SPP(nn.Module):
    # Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729
    def __init__(self, c1, c2, k=(5, 9, 13)):
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
 
    def forward(self, x):
        x = self.cv1(x)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning
            return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))

YoloV5中的SPP类定义了一个空间金字塔池化（Spatial Pyramid Pooling, SPP）层，这是一种常用于卷积神经网络的模块，旨在改善模型对于不同尺寸输入的适应能力。SPP层最早由Kaiming He等人在其2014年的论文《Spatial Pyramid Pooling in Deep Convolutional Networks for Visual Recognition》中提出。

SPP的基本原理是使用不同尺寸的最大池化（max pooling）操作来处理特征图，这帮助网络能够捕获到不同尺度下的特征。具体在这段代码中，SPP层的结构如下：

Conv(c1, c_, 1, 1): 首先使用一个1x1的卷积核进行特征压缩，将输入通道c1减半到隐藏通道c_。
nn.ModuleList([...]): 这段代码创建了一个模块列表，包含了多个不同大小的最大池化层。每个池化层的核尺寸由k参数给出，默认值为(5, 9, 13)。池化层用于在不同的感受野（kernel size）上对输入x进行池化操作，stride设为1以保持特征图的尺寸，padding设置为kernel_size // 2以确保边界的特征同样被考虑。
self.cv2: 之后，使用一个1x1的卷积将所有的池化输出与原始的特征图x拼接起来（这是在forward函数中执行的），并将拼接后的特征图压缩为输出通道c2。
forward(self, x): 在前向传播过程中，首先对输入特征x执行cv1层定义的1x1卷积，然后将这个结果和通过不同尺寸池化层处理后的特征图进行拼接。拼接后的结果通过cv2定义的另一个1x1卷积完成特征融合。为了避免PyTorch 1.9.0版本中的一个特定警告，代码中使用了一个warnings上下文管理器。

通过使用SPP层，网络能够收集和融合不同尺度的上下文信息，这可以提高模型对尺度变化的鲁棒性，并且由于SPP允许任意大小的输入，因而模型在处理不同尺寸的输入图像时更为灵活。这些特性使得SPP层特别适用于检测任务，尤其是在对象大小和图像尺寸可能有很大变化的情况下。在YOLOv5这样的目标检测模型中，这有助于提高模型对不同尺寸目标的检测能力。

8.Contract

class Contract(nn.Module):
    # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
    def __init__(self, gain=2):
        super().__init__()
        self.gain = gain
 
    def forward(self, x):
        b, c, h, w = x.size()  # assert (h / s == 0) and (W / s == 0), 'Indivisible gain'
        s = self.gain
        x = x.view(b, c, h // s, s, w // s, s)  # x(1,64,40,2,40,2)
     	# permute：改变tensor的维度顺序
        x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # x(1,2,2,64,40,40)
        # .view：改变tensor的维度
        return x.view(b, c * s * s, h // s, w // s)  # x(1,256,40,40)

Contract 类在 YOLOv5 的代码中定义了一个自定义的神经网络模块，其功能是改变张量 x 的尺寸，通过降低高度（h）和宽度（w），将这部分的信息"收缩"到通道（c）里。

这个过程起到下采样的作用，但它与常规的池化或卷积操作不同，因为它是通过重新排列张量的维度来实现的。详细过程如下：

初始化时，设置 gain 值（默认为 2），这个值指定了在高度和宽度上的缩减因子。
在前向传播（forward 函数）中：
- 首先获取张量 x 的大小（batch size b, channels c, height h, width w）。
- 设置缩减因子 s 为 gain 的值。
- x 被重新视图（view）为一个具有新维度顺序的张量，通过将高度和宽度分别除以 s 并将这些"收缩"到了新的维度中，比如从 (b, c, h, w) 重塑为 (b, c, h // s, s, w // s, s)。
- 使用 permute 函数来改变张量的维度顺序，以将收缩后的空间维度与通道维度交错起来。
- 最终使用 view 将张量重塑为 (b, c * s * s, h // s, w // s)，以便保持总的数据量不变，但空间分辨率减少。

这个操作实质上是对每个小块（大小由 gain 定义）在宽度和高度上进行重新排列并移动这些小块信息到通道维度上。结果是，对于每个 (s, s) 的空间区块，局部空间信息被有效"压缩"到了深度通道中。这种方法带来的可能好处包括减少计算量（由于空间分辨率减少）以及可能增加特征的混合，这可能有助于神经网络的学习效率。

9.Concat

class Concat(nn.Module):
    # Concatenate a list of tensors along dimension
    def __init__(self, dimension=1):
        super().__init__()
        self.d = dimension
 
    def forward(self, x):
        return torch.cat(x, self.d)

在 YOLOv5 的代码中，Concat 类定义了一个非常简单且常用的神经网络模块，其作用是将一个列表中的多个张量沿指定的维度拼接起来。

这个操作在构建深度学习模型，尤其是卷积神经网络中很有用，因为它允许将不同的特征层或者来自之前层的特征图拼接到一起。这就为网络提供了一种机制来组合来自不同处理流程的信息。举例来说，在特征金字塔网络（Feature Pyramid Networks, FPNs）或者 U-Net 结构中，拼接操作用于结合来自网络不同级别的特征，以实现多尺度特征融合。

具体到 Concat 类的代码功能：

在初始化(__init__)过程中，Concat 类接收一个参数 dimension，默认值为 1。这个参数指定了拼接发生在哪一个维度上。在一个具有形状 [batch_size, channels, height, width] 的四维张量中，通常是沿通道维度拼接，即 dimension=1。
在前向传播 (forward) 函数中，它接收一个张量列表 x，并调用 PyTorch 的 torch.cat 函数，沿着 self.d 指定的维度将这些张量拼接起来。

这种方法经常用来：

合并来自相同特征图层次的不同分支的输出。
在上采样（增加分辨率）后与前面某个下采样层的输出进行拼接，这是 U-Net 架构的典型做法。

例如，在 YOLOv5 的模型中，可能会从不同的层面提取特征，然后将这些不同层的输出拼接起来，以此来保留高分辨率的精细特征和低分辨率但是语义信息更丰富的特征。这种特征拼接有助于提升模型检测不同尺度物体的能力。

二、重要类

1.非极大值抑制（NMS）

class NMS(nn.Module):
    """在yolo.py中Model类的nms函数中使用
    NMS非极大值抑制 Non-Maximum Suppression (NMS) module
    给模型model封装nms  增加模型的扩展功能  但是我们一般不用 一般是在前向推理结束后再调用non_max_suppression函数
    """
    conf = 0.25     # 置信度阈值              confidence threshold
    iou = 0.45      # iou阈值                IoU threshold
    classes = None  # 是否nms后只保留特定的类别 (optional list) filter by class
    max_det = 1000  # 每张图片的最大目标个数    maximum number of detections per image
    def __init__(self):
        super(NMS, self).__init__()
    def forward(self, x):
        """
        :params x[0]: [batch, num_anchors(3个yolo预测层), (x+y+w+h+1+num_classes)]
        直接调用的是general.py中的non_max_suppression函数给model扩展nms功能
        """
        return non_max_suppression(x[0], self.conf, iou_thres=self.iou, classes=self.classes, max_det=self.max_det)

在 YOLOv5 架构中，NMS 类是为了在模型内实现非极大值抑制（Non-Maximum Suppression，NMS）功能的一个封装。这是一种在目标检测任务中经常使用的后处理步骤，用来减少相邻的、多余的检测框，并保留最佳的单一检测结果。

非极大值抑制的流程大致如下：

确定所有检测框的置信度，并丢弃那些置信度较低（低于某个阈值）的检测框。
在剩下的检测框中，选择置信度最高的。
计算这个检测框与其他所有检测框的 IOU（交并比），并丢弃那些 IOU 高于特定阈值的检测框。
重复步骤 2 和 3 直到所有检测框都经过处理。

NMS 类允许定义相关的参数：

conf: 置信度阈值，用来筛选掉那些目标检测置信度低于此阈值的预测。
iou: IOU 阈值，用于在 NMS 过程中判断检测框之间的重叠程度。如果检测框与最高得分检测框的 IOU 超过此值，则被认为是重叠，并且将被抑制。
classes: 是否只保留 NMS 后特定类别的检测结果。如果设置了特定的类别列表，则只有列表中的类别会被保留。
max_det: 每张图片最大可保留的检测个数，防止对于某些图片预测出大量的检测框。

NMS模块里的 forward 方法实际上调用了 non_max_suppression 函数，这是在 YOLOv5 的工具库 general.py 中一个已经定义好的函数。NMS 类的存在允许在构建模型时将 NMS 作为一个层直接集成进去，而不是作为预测后的附加步骤。这可以让模型变得更加模块化，并可能在部署时简化步骤。

不过，通常情况下，由于在训练过程中不需要进行 NMS，所以这一步骤往往是在模型前向传递结束后，在得到预测框之后单独进行的。因此在实际应用中，更常在推理后单独调用 NMS 函数，而不是嵌入到模型结构中。这样做的好处是可以根据实际需求灵活地调整 NMS 参数，例如在不同的测试场景下可能需要不同的置信度和 IOU 阈值。

2.AutoShape

class AutoShape(nn.Module):
    # YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
    # YOLOv5模型包装器，用于传递 cv2/np/PIL/torch 输入，
    # 包括预处理（preprocessing）, 推理（inference） and NMS
    conf = 0.25  # NMS confidence threshold
    iou = 0.45   # NMS IoU threshold
    agnostic = False  # NMS class-agnostic
    multi_label = False  # NMS multiple labels per box
    classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
    max_det = 1000  # maximum number of detections per image
    amp = False     # Automatic Mixed Precision (AMP) inference
 
    def __init__(self, model, verbose=True):
        super().__init__()
        if verbose:
            LOGGER.info('Adding AutoShape... ')
        copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=())  # copy attributes
        self.dmb = isinstance(model, DetectMultiBackend)  # DetectMultiBackend() instance
        self.pt = not self.dmb or model.pt  # PyTorch model
        # 开启验证模式
        self.model = model.eval()
        if self.pt:
            m = self.model.model.model[-1] if self.dmb else self.model.model[-1]  # Detect()
            m.inplace = False  # Detect.inplace=False for safe multithread inference
            m.export = True  # do not output loss values
 
    def _apply(self, fn):
        # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
        self = super()._apply(fn)
        if self.pt:
            m = self.model.model.model[-1] if self.dmb else self.model.model[-1]  # Detect()
            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))
        return self
 
    @smart_inference_mode()
    def forward(self, ims, size=640, augment=False, profile=False):
        # Inference from various sources. For size(height=640, width=1280), RGB images example inputs are:
        #   file:        ims = 'data/images/zidane.jpg'  # str or PosixPath
        #   URI:             = 'https://ultralytics.com/images/zidane.jpg'
        #   OpenCV:          = cv2.imread('image.jpg')[:,:,::-1]  # HWC BGR to RGB x(640,1280,3)
        #   PIL:             = Image.open('image.jpg') or ImageGrab.grab()  # HWC x(640,1280,3)
        #   numpy:           = np.zeros((640,1280,3))  # HWC
        #   torch:           = torch.zeros(16,3,320,640)  # BCHW (scaled to size=640, 0-1 values)
        #   multiple:        = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...]  # list of images
 
        dt = (Profile(), Profile(), Profile())
        with dt[0]:
            if isinstance(size, int):  # expand
                size = (size, size)
            p = next(self.model.parameters()) if self.pt else torch.empty(1, device=self.model.device)  # param
            autocast = self.amp and (p.device.type != 'cpu')  # Automatic Mixed Precision (AMP) inference
            # 图片如果是tensor格式 说明是预处理过的, 
            # 直接正常进行前向推理即可 nms在推理结束进行(函数外写)
            if isinstance(ims, torch.Tensor):  # torch
                with amp.autocast(autocast):
                    return self.model(ims.to(p.device).type_as(p), augment=augment)  # inference
 
            # Pre-process
            n, ims = (len(ims), list(ims)) if isinstance(ims, (list, tuple)) else (1, [ims])  # number, list of images
            shape0, shape1, files = [], [], []  # image and inference shapes, filenames
            for i, im in enumerate(ims):
                f = f'image{i}'  # filename
                if isinstance(im, (str, Path)):  # filename or uri
                    im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im), im
                    im = np.asarray(exif_transpose(im))
                elif isinstance(im, Image.Image):  # PIL Image
                    im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename', f) or f
                files.append(Path(f).with_suffix('.jpg').name)
                if im.shape[0] < 5:  # image in CHW
                    im = im.transpose((1, 2, 0))  # reverse dataloader .transpose(2, 0, 1)
                im = im[..., :3] if im.ndim == 3 else cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)  # enforce 3ch input
                s = im.shape[:2]  # HWC
                shape0.append(s)  # image shape
                g = max(size) / max(s)  # gain
                shape1.append([int(y * g) for y in s])
                ims[i] = im if im.data.contiguous else np.ascontiguousarray(im)  # update
            shape1 = [make_divisible(x, self.stride) for x in np.array(shape1).max(0)]  # inf shape
            x = [letterbox(im, shape1, auto=False)[0] for im in ims]  # pad
            x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2)))  # stack and BHWC to BCHW
            x = torch.from_numpy(x).to(p.device).type_as(p) / 255  # uint8 to fp16/32
 
        with amp.autocast(autocast):
            # Inference
            with dt[1]:
                y = self.model(x, augment=augment)  # forward
            # Post-process
            with dt[2]:
                y = non_max_suppression(y if self.dmb else y[0],
                                        self.conf,
                                        self.iou,
                                        self.classes,
                                        self.agnostic,
                                        self.multi_label,
                                        max_det=self.max_det)  # NMS
                for i in range(n):
                    scale_boxes(shape1, y[i][:, :4], shape0[i])
            return Detections(ims, y, files, dt, self.names, x.shape)

AutoShape 类在 YOLOv5 模型中是一个通用模型包装器，用于简化从不同来源的输入数据进行预处理、模型推理及非极大值抑制（NMS）的整个流程。这个类设计的主要目的是为了能够灵活地处理各种输入格式，例如来自 OpenCV、Numpy、PIL 或 Torch 的图像数据。

以下是 AutoShape 类的主要作用和流程：

自动化输入预处理: 不同的图像来源可能需要不同的加载和预处理步骤。例如，来自文件路径的图像、来自网址的图像、或是已经加载到内存中的图像数据类型会有所不同。AutoShape 类自动处理这些不同的输入，确保它们都被正确加载并转换成模型需要的格式。
简化模型推理（Inference）: 该类通过封装模型，允许模型直接处理不同的输入格式。接收到输入后，会进行必要的尺寸调整、归一化等预处理步骤，随后执行模型前向推理。
封装后处理: 在模型推理后，AutoShape 类还会自动对输出结果进行非极大值抑制(NMS)，这是目标检测后处理的关键步骤，用于清理重叠的检测框并保留最佳的检测结果。
灵活性: 在推理时，AutoShape 类实现了对自动混合精度推理（AMP）的支持，这可以加速模型的推理速度，并降低硬件（如GPU）的内存使用，而不会显著降低推理质量。
配置参数: 类似 NMS 类，AutoShape 也有一些可配置的属性，如置信度阈值、IOU 阈值、类别过滤设置、每图像最大检测数等。这些设置在进行 NMS 处理时使用。

在 forward 方法中：

它会检查输入数据的类型，并根据类型处理图像数据（加载图像、预处理图像）。
图像会被调整至合适的尺寸，并且会进行归一化和格式化以满足模型输入要求。
应用模型推理，并使用设定的参数完成 NMS 处理。
通过 Detections 对象组织输出格式，包括推理结果以及预处理和推理所用的时间分析。

AutoShape 类大大减少了实施 YOLOv5 模型推理时所需编写的代码量，并且为使用者提供了一种简单直接处理不同图像来源的方法。

3.Detections

 
class Detections:
    # YOLOv5 detections class for inference results
    # YOLOv5推理结果检测类
    def __init__(self, ims, pred, files, times=(0, 0, 0), names=None, shape=None):
        super().__init__()
        d = pred[0].device  # device
        gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in ims]  # normalizations
        self.ims = ims  # list of images as numpy arrays
        self.pred = pred  # list of tensors pred[0] = (xyxy, conf, cls)
        self.names = names  # class names
        self.files = files  # image filenames
        self.times = times  # profiling times
        self.xyxy = pred  # xyxy pixels
        self.xywh = [xyxy2xywh(x) for x in pred]  # xywh pixels
        self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)]  # xyxy normalized
        self.xywhn = [x / g for x, g in zip(self.xywh, gn)]  # xywh normalized
        self.n = len(self.pred)  # number of images (batch size)
        self.t = tuple(x.t / self.n * 1E3 for x in times)  # timestamps (ms)
        self.s = tuple(shape)  # inference BCHW shape
 
    def _run(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path('')):
        s, crops = '', []
        for i, (im, pred) in enumerate(zip(self.ims, self.pred)):
            s += f'\nimage {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '  # string
            if pred.shape[0]:
                for c in pred[:, -1].unique():
                    n = (pred[:, -1] == c).sum()  # detections per class
                    s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "  # add to string
                s = s.rstrip(', ')
                if show or save or render or crop:
                    annotator = Annotator(im, example=str(self.names))
                    for *box, conf, cls in reversed(pred):  # xyxy, confidence, class
                        label = f'{self.names[int(cls)]} {conf:.2f}'
                        if crop:
                            file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
                            crops.append({
                                'box': box,
                                'conf': conf,
                                'cls': cls,
                                'label': label,
                                'im': save_one_box(box, im, file=file, save=save)})
                        else:  # all others
                            annotator.box_label(box, label if labels else '', color=colors(cls))
                    im = annotator.im
            else:
                s += '(no detections)'
 
            im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im  # from np
            if show:
                display(im) if is_notebook() else im.show(self.files[i])
            if save:
                f = self.files[i]
                im.save(save_dir / f)  # save
                if i == self.n - 1:
                    LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
            if render:
                self.ims[i] = np.asarray(im)
        if pprint:
            s = s.lstrip('\n')
            return f'{s}\nSpeed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {self.s}' % self.t
        if crop:
            if save:
                LOGGER.info(f'Saved results to {save_dir}\n')
            return crops
 
    @TryExcept('Showing images is not supported in this environment')
    def show(self, labels=True):
        self._run(show=True, labels=labels)  # show results
 
    def save(self, labels=True, save_dir='runs/detect/exp', exist_ok=False):
        save_dir = increment_path(save_dir, exist_ok, mkdir=True)  # increment save_dir
        self._run(save=True, labels=labels, save_dir=save_dir)  # save results
 
    def crop(self, save=True, save_dir='runs/detect/exp', exist_ok=False):
        save_dir = increment_path(save_dir, exist_ok, mkdir=True) if save else None
        return self._run(crop=True, save=save, save_dir=save_dir)  # crop results
 
    def render(self, labels=True):
        self._run(render=True, labels=labels)  # render results
        return self.ims
 
    def pandas(self):
        # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])
        new = copy(self)  # return copy
        ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name'  # xyxy columns
        cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name'  # xywh columns
        for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):
            a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)]  # update
            setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])
        return new
 
    def tolist(self):
        # return a list of Detections objects, i.e. 'for result in results.tolist():'
        r = range(self.n)  # iterable
        x = [Detections([self.ims[i]], [self.pred[i]], [self.files[i]], self.times, self.names, self.s) for i in r]
        # for d in x:
        #    for k in ['ims', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
        #        setattr(d, k, getattr(d, k)[0])  # pop out of list
        return x
 
    def print(self):
        LOGGER.info(self.__str__())
 
    def __len__(self):  # override len(results)
        return self.n
 
    def __str__(self):  # override print(results)
        return self._run(pprint=True)  # print results
 
    def __repr__(self):
        return f'YOLOv5 {self.__class__} instance\n' + self.__str__()

4.Classify

class Classify(nn.Module):
    # YOLOv5 classification head, i.e. x(b,c1,20,20) to x(b,c2)
    def __init__(self,
                 c1,
                 c2,
                 k=1,
                 s=1,
                 p=None,
                 g=1,
                 dropout_p=0.0):  # ch_in, ch_out, kernel, stride, padding, groups, dropout probability
        super().__init__()
        c_ = 1280  # efficientnet_b0 size
        self.conv = Conv(c1, c_, k, s, autopad(k, p), g)
        self.pool = nn.AdaptiveAvgPool2d(1)  # to x(b,c_,1,1)
        self.drop = nn.Dropout(p=dropout_p, inplace=True)
        self.linear = nn.Linear(c_, c2)  # to x(b,c2)
 
    def forward(self, x):
        if isinstance(x, list):
            x = torch.cat(x, 1)
        return self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))

三、C3修改

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.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
 
    def forward(self, x):
        return torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1)