读yolov5代码记录

因为这两天很多东西要用到yolov5,所以打算近期把代码整个过一遍,代码


utils

utils:这个文件夹包含了一些工具函数,如数据加载、边界框处理等。从YOLOv5中的数据处理和辅助函数读起有助于后续理解模型。

文件结构

在这里插入图片描述
从文件夹名来看,这些文件夹不重要,对本地运行代码没有什么作用,应该属部署到指定云服务器的一些操作,直接看代码!

init.py

init.py:这个文件没什么作用,四个函数一个类,第一个函数是针对Windows的emoji的编码转换;线程调度相关的两个函数,略;notebook初始化相关的函数;一个类是捕获异常并打印。

activations.py

activations.py:这段代码包含了一些针对深度学习神经网络的激活函数。

  1. SiLU:将输入向量逐元素乘以其Sigmoid映射值,引入非线性。公式:f(x) = x * sigmoid(x)。
  2. Hardswish:逐元素版本的SiLU(Hard-SiLU),端点区域具有平滑特性。公式:f(x) = x * hardtanh(x + 3) / 6,其中硬双曲正切hardtanh(x) = min(max(x, lower_limit), upper_limit)。
  3. Mish:具有高斯平滑特性的非线性激活函数。公式:f(x) = x * tanh(softplus(x)),其中softplus(x) = log(1 + e^x)。
  4. MemoryEfficientMish:在计算梯度的内存使用效率方面优化的Mish激活函数。计算方式与Mish相同。
  5. FReLU:在ReLU的基础上增加卷积、批量归一化操作以提高性能。公式:f(x) = max(x, α∗Conv2d(x)),其中α为可学习参数。
  6. AconC:具有可学习参数的激活函数,允许自定义激活。公式:(p1 * x - p2 * x) * sigmoid(beta * (p1 * x - p2 * x)) + p2 * x,其中p1, p2, beta为可学习参数。
  7. MetaAconC:与AconC类似的激活函数,通过小型神经网络生成可学习参数。公式与 AconC 相似,但将固定的beta参数替换为由神经网络生成的参数。

augmentations.py

augmentations.py:这段代码主要是为yolo提供图像增强的操作。
Albumentations 类,用于封装可选的 Albumentations 图像增强。如果库被安装并正确导入,该类提供的功能包括:RandomResizedCrop(随机重新缩放和裁剪)、Blur(模糊)、MedianBlur(邻近中值模糊)、ToGray(转灰度)、CLAHE(对比度受限的自适应直方图均衡化提高局部对比度)、RandomBrightnessContrast 和 RandomGamma、ImageCompression(图像压缩模拟压缩失真 )。

autoanchor.py

``


models

models: 这个文件夹包含YOLOv5的模型定义和训练相关的代码。从这里开始查看YOLOv5的核心模型结构。

文件结构

models部分的文件结构

hub文件夹:

anchors.yaml:COCO数据集下,不同输入图像和目标大小的锚框长宽

# P6-640:  thr=0.25: 0.9964 BPR, 5.54 anchors past thr, n=12, img_size=640, metric_all=0.281/0.716-mean/best, past_thr=0.469-mean: 9,11,  21,19,  17,41,  43,32,  39,70,  86,64,  65,131,  134,130,  120,265,  282,180,  247,354,  512,387

注释是关于锚框参数(anchors_p6_640)的一段描述,其中提供了关于该锚框参数组的一些性能指标和统计信息。如下逐个解释这些字段:

  • P6-640: 表明这些数据是针对anchors_p6_640设计的,覆盖了P6/64特征图,图像尺 寸(img_size)为640。
  • thr=0.25: 设置了一个阈值,用于评估锚框的性能。通常,这意味着只有匹配度超过0.25的锚框才会被认为是“好”的锚框。
  • 0.9964 BPR: BPR(Best Possible Recall)是衡量锚框参数性能的指标之一,表示最佳可能召回率。这里,该参数组的BPR值为0.9964,意味着在特定阈值(0.25)下,锚框的召回率接近1,这是一个很高的值。
  • 5.54 anchors past thr: 表示平均每个ground truth目标有5.54个锚框超过了设定的阈值(0.25)。
  • n=12: 是用于生成这些统计数据的锚框数量, 这里是12个锚框。
  • metric_all=0.281/0.716-mean/best: 这是针对所有锚框设计的评估指标(可能是IoU等)的平均值(0.281)和最佳值(0.716)。
  • past_thr=0.469-mean: 过滤阈值(0.25)后的锚框评估指标的平均值为0.469。 之后的数字序列,例如9,11, 21,19, 17,41, 43,32, 39,70, 86,64, 65,131, 134,130, 120,265, 282,180,
    247,354, 512,387,是具体锚框尺寸(宽度和高度)的列表。这些数字成对出现,表示锚框的宽度和高度。

yolov3-spp.yaml:该文件以及类似格式的文件给出了具体yolo的网络结构,主要是骨干网络和检测头的结构。

下面是YOLOv5配置文件中各个参数的解释:

  • nc: 类别数量,这里设置为80,即80个不同的目标类型。
  • depth_multiple与width_multiple: 用于调整模型规模,分别表示模型深度(层数)倍数和每层通道数倍数。这里均设置为1.0,表示模型采用标准配置。
  • anchors: 锚框大小,这里分为3组,分别对应特征图P3/8、P4/16和P5/32。锚框尺寸有 width 和 height 成对组成。

接下来是两个主要部分:backbone 和 head。

  • backbone 部分是针对YOLOv5 v6.0的骨干网络结构的描述。

    1. -1表示从上一个模块或层的输出获取输入。 number指定要重复的模块数量(可用于堆叠多个相同结构的模块)。
    2. module表示要使用的模块类型(例如 Conv,C3 等)。
    3. args是一个包含模块参数的列表。如输入通道数量,卷积核大小等。
  • head部分使用类似的描述方式,包含模块类型和参数。

最后,使用 [17, 20, 23], 1, Detect, [nc, anchors] 表示在P3, P4, P5特征图上应用目标检测,其中nc表示目标类别数量,anchors表示设置的锚框大小。

segment文件夹

  • 找到这部分的作用
    从文件名来说,应该是用yolov5进行分割任务的一些参数...

common.py

common.py文件

  • def autopad(k, p=None, d=1)

输入是卷积核大小k,填充p,以及展开d。这个函数的主要目的是给定卷积核,填充和展开的设定,为实现’same’卷积计算而得出一个理想的填充

  • class Conv(nn.Module)
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 类是一个卷积层,包括一个卷积层、一个BN层和一个激活函数。有两种 forward 函数,forward 和 forward_fuse,分别表示经过 batch normalization 层与否

  • class DWConv(Conv)
class DWConv(Conv):
    # Depth-wise convolution
    def __init__(self, c1, c2, k=1, s=1, d=1, act=True):  # ch_in, ch_out, kernel, stride, dilation, activation
        super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), d=d, act=act)

DWConv ("Depthwise Separable Convolution "),即深度可分离卷积,是一种特殊类型的卷积,它比标准卷积运算更为高效。深度卷积的操作是在每个输入通道上单独应用卷积滤波器,然后再对结果进行聚合,这大大减少了计算量。

	假设我们的输入是一个形状为[10, 10, 32]的特征图,我们想让输出的通道数为64卷积核的大小为[3, 3]:
	- 标准卷积:我们需要的卷积核的大小为[3, 3, 32, 64],所以总的参数数量是3*3*32*64 = 18432个。加上卷积需要的计算量是10*10*3*3*32*64 = 1843200
	- 深度可分离卷积:我们首先进行深度卷积,需要32个[3, 3, 1]卷积核,参数数量是3*3*32 = 288然后进行1x1卷积,加入输入通道之间的信息交换,卷积核大小是[1, 1, 32, 64],参数数量是32*64=2048。所以深度可分离卷积的总参数数量是288+2048=2336。
	与标准卷积相比,参数数量大大减少了,计算量也大大减少了。

从代码来,这应该是一个分组卷积,就是将输入通道分为多个组分别卷积,参数量为普通卷积参数量/组数

  • class DWConvTranspose2d(nn.ConvTranspose2d)

DWConvTranspose2d 与 DWConv 类似,但它执行的是转置卷积(反卷积),因为卷积操作中的值可以转换为矩阵操作,举个例子,44的图,经过33卷积核,变成22,这个22中的每一个值都是卷积核中元素跟图中元素相乘相加得到的,换言之就是416的矩阵与161的矩阵相乘,得到4*1的矩阵,这种操作可以增加输入数据的空间大小(高度和宽度),关于反卷积这里附一个链接:反卷积

  • class TransformerLayer(nn.Module)
class TransformerLayer(nn.Module):
    # Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
    def __init__(self, c, num_heads):
        super().__init__()
        self.q = nn.Linear(c, c, bias=False)
        self.k = nn.Linear(c, c, bias=False)
        self.v = nn.Linear(c, c, bias=False)
        self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
        self.fc1 = nn.Linear(c, c, bias=False)
        self.fc2 = nn.Linear(c, c, bias=False)

    def forward(self, x):
        x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
        x = self.fc2(self.fc1(x)) + x
        return x

在__init__函数中,参数c代表输入和输出的特征维度个数,num_heads 代表指定的attention heads 的数量。q,k,v是线性层用于生成“查询(query)”,“键(key)”和“值(value)”。
在前向传播(forward)函数中,输入 x 首先经过 self.q、self.k和 self.v,然后传入 self.ma(多头注意力)。得到的输出再与原始 x 相加形成一种残差连接。然后,结果通过两个全连接层,并再次形成残差连接。
多头注意力
在多头注意力中,"多头"意味着模型不仅在一个表示空间(或者说一个“注意力头”)上学习输入的表表示,而是在多个表示空间上同时进行,每个表示空间都有自己的学习参数。这为模型提供了更多的能力去获取不同类型的信息和潜在模式。
例如,当处理一个句子时,一个"头"可以专注于词语的语法关系,另一个"头"则可以看重词义之间的相关性。这样,多头注意力就能捕捉输入中更丰富、更复杂的特征。

  • class TransformerBlock(nn.Module)
class TransformerBlock(nn.Module):
    # Vision Transformer https://arxiv.org/abs/2010.11929
    def __init__(self, c1, c2, num_heads, num_layers):
        super().__init__()
        self.conv = None
        if c1 != c2:
            self.conv = Conv(c1, c2)
        self.linear = nn.Linear(c2, c2)  # learnable position embedding
        self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers)))
        self.c2 = c2

    def forward(self, x):
        if self.conv is not None:
            x = self.conv(x)
        b, _, w, h = x.shape
        p = x.flatten(2).permute(2, 0, 1)
        return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)

在__init__函数中,c1和c2 代表输入和输出的特征维度个数,num_heads代表了注意力头部的数量,num_layers代表Transformer层的数量。
Conv(c1, c2)是前文提到的卷积层,包括卷积、BN以及激活函数,用于学习位置嵌入。self.tr 则是由多个TransformerLayer组成的序列。
在前向传播(forward)函数中,如果存在self.conv,则对输入 x 进行卷积变换。然后将 x 转换为2D形状,并且调换维度使其适合传入 transformer。p 和 self.linear( p) 相加后被送入指定层数的transform序列。最后,调换维度并将其重新塑造为输入 tensor 的形状。

  • class Bottleneck(nn.Module)
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
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_, c2, 3, 1, g=g)
        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模块根据self.add的值操作,如果为True,则进行两层卷积和元素的相加操作,相当于ResNet的shortcut连接;如果为False,则只保留两层卷积操作的结果。

  • class BottleneckCSP(nn.Module)
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))))

Cross Stage Partial Networks (CSPNet)的Bottleneck模块。CSPNet 是通过在特征图上构建跨层连接的方式来改进神经网络的构建方式,避免特征的重复聚合和信息冗余,从而实现更好的性能。
模型首先经过 cv1 卷积层,并通过 m 函数的剩余连接进行处理,然后再针对该结果再次应用 cv3 卷积层,将输出保存为 y1。与此同时,输入 x 直接经过 cv2 卷积后的结果为 y2。然后,y1 和 y2 在通道维度进行拼接,并且这个合并的输出经过批标准化 (bn) 和激活函数 (act) 后再通过 cv4 卷积层进行处理,给出最后结果。也就是说, cv2 和 cv3 和它们之间的连接形成"shortcut"或"skip"连接。

  • class CrossConv(nn.Module)
class CrossConv(nn.Module):
    # Cross Convolution Downsample
    def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
        # ch_in, ch_out, kernel, stride, groups, expansion, shortcut
        super().__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, (1, k), (1, s))
        self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
        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))

CrossConv是一个包含了两个卷积层的模块。第一个卷积层是一个(1,k)大小的卷积核,第二个卷积层是一个(k,1)大小的卷积核,它们共同作用于输入的数据。如果输入和输出的通道数一致,且shortcut为True,则它会将原始输入与卷积结果相加返回,这是一个残差连接,可以帮助解决深度网络模型训练过程中的退化问题。
这种类型的交叉卷积操作模式可以帮助网络在保持准确性的同时节省计算资源,因为1k和k1的两个小卷积相比一个k*k的大卷积来说计算量小很多。对神经网络模型进行此类操作可以增加其效率

  • class C3(nn.Module)
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是CSP(Cross Stage Partial networks)Bottleneck模块。CSPNet是一种新的深度学习网络结构,其主要贡献是通过设计跨阶段的部分连接结构,将网络的特征图分割为两部分进行运算,从而降低网络的内存消耗。在这个模块中,输入x被分为两个路径,一部分经过cv1卷积以及多个Bottleneck模块,另一部分经过cv2卷积,然后这两部分特征图在通道维度上被拼接起来,最后通过cv3卷积输出。

  • class C3x(C3)、 C3TR(C3)、 C3SPP(C3)、C3Ghost(C3)
class C3x(C3):
    # C3 module with cross-convolutions
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
        super().__init__(c1, c2, n, shortcut, g, e)
        c_ = int(c2 * e)
        self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)))


class C3TR(C3):
    # C3 module with TransformerBlock()
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
        super().__init__(c1, c2, n, shortcut, g, e)
        c_ = int(c2 * e)
        self.m = TransformerBlock(c_, c_, 4, n)


class C3SPP(C3):
    # C3 module with SPP()
    def __init__(self, c1, c2, k=(5, 9, 13), n=1, shortcut=True, g=1, e=0.5):
        super().__init__(c1, c2, n, shortcut, g, e)
        c_ = int(c2 * e)
        self.m = SPP(c_, c_, k)


class C3Ghost(C3):
    # C3 module with GhostBottleneck()
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
        super().__init__(c1, c2, n, shortcut, g, e)
        c_ = int(c2 * e)  # hidden channels
        self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))

C3x、C3TR,C3SPP 和 C3Ghost 使用一些模块(CrossConv、TransformerBlock, SPP, GhostBottleneck)对C3 (CSP Bottleneck with 3 convolutions) 进行扩展。基础构造函数中的参数包括输入通道数目 (c1)、输出通道数目 (c2)、重复模块数目 (n)、是否需要shortcut操作 (shortcut)、分组卷积的组数目 (g),未完全连接 (Unconnected, e) 的比例。说白了就是一些卷积模块的叠加。

  • class SPP(nn.Module)
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))

SPP模块的作用就是处理任意尺度的输入,并且能够从各种尺度和位置上提取特征,因为在目标检测任务中目标的尺寸大小未知,使用SSP取代之前的滑动窗口可以减少计算量,同时避免漏掉一些全局信息。
具体实现方式为

  • class SPPF(nn.Module)
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))

Spatial Pyramid Pooling - Fast,快速空间金字塔池化,传统的SPP(Spatial Pyramid Pooling)操作在一个空间尺度上(单一图片)获取多尺度(即代码中的5,9,13)特征,而SPFF(Spatial Pyramid Pooling - Fast)则利用同一个核多次池化,获得了输入图片的不同粒度特征。这种策略在保持传统SPP特性的同时,降低了模型复杂性

  • class Focus(nn.Module)
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))

这部分从代码来看有点抽象,简单来说就是隔项取,行和列分别隔项取,就会变成四个矩阵:

拼接后张量的形状为 (b,4c,w/2,h/2),这意味着通过从输入张量的不同子区域抽取信息,扩大了通道数(即深度)。

  • class GhostConv(nn.Module)
class GhostConv(nn.Module):
    # Ghost Convolution https://github.com/huawei-noah/ghostnet
    def __init__(self, c1, c2, k=1, s=1, g=1, act=True):  # ch_in, ch_out, kernel, stride, groups
        super().__init__()
        c_ = c2 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, k, s, None, g, act=act)
        self.cv2 = Conv(c_, c_, 5, 1, None, c_, act=act)

    def forward(self, x):
        y = self.cv1(x)
        return torch.cat((y, self.cv2(y)), 1)

通过廉价操作生成更多的特征图。基于一组原始的特征图,作者应用一系列线性变换,以很小的代价生成许多能从原始特征发掘所需信息的“Ghost”特征图。这是因为不同层次卷积输出的特征图可能高度相似,基本只要进行简单的线性变换就能得到,而不需要进行复杂的非线性变换得到。
具体来说,就是将原特征图降维之后对小特征图卷积,之后将小特征图与卷积结果线性拼接,假设输入特征图的shape为[28,28,6],首先对输入特征图使用11卷积下降通道数,shape变为[28,28,3];再使用33深度卷积对每个通道特征图提取特征,shape为[28,28,3],可以看作是经过前一层的一系列线性变换得到的;最后将两次卷积的输出特征图在通道维度上堆叠,shape 变为 [28,28,6];
ghost卷积

  • class GhostBottleneck(nn.Module)

应用深度可分离卷积和Ghost卷积,类似残差网络结构,具体如下图:
在这里插入图片描述

  • class Contract(nn.Module)
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)
        x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # x(1,2,2,64,40,40)
        return x.view(b, c * s * s, h // s, w // s)  # x(1,256,40,40)

这个模块会把输入图像的空间维度(宽和高)压缩,同时增加通道维度的大小。具体操作就是分离维度,重新排列,再次组合,举个例子,它可以把形状为 (1, 64, 80, 80) 的输入变为(1,64,40,2,40,2),再变为 (1,2,2,64,40,40),最后view函数得到(1, 256, 40, 40) 的输出,也就是通过收缩空间维度将通道数从 64 增加到了 256。

  • class Expand(nn.Module)
class Expand(nn.Module):
    # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)
    def __init__(self, gain=2):
        super().__init__()
        self.gain = gain

    def forward(self, x):
        b, c, h, w = x.size()  # assert C / s ** 2 == 0, 'Indivisible gain'
        s = self.gain
        x = x.view(b, s, s, c // s ** 2, h, w)  # x(1,2,2,16,80,80)
        x = x.permute(0, 3, 4, 1, 5, 2).contiguous()  # x(1,16,80,2,80,2)
        return x.view(b, c // s ** 2, h * s, w * s)  # x(1,16,160,160)

这个模块和Contract模块相反,但是操作实现类似

  • class Concat(nn.Module)
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)

这个模块比较简单,就是将向量在指定维度上拼接,比如两个22矩阵,在第一个维度上拼接就是42矩阵,在第二个维度上拼接就是2*4矩阵

  • class DetectMultiBackend(nn.Module)

这是一个用多种后端(如 PyTorch、TorchScript、ONNX Runtime、OpenVINO等)进行推断的类,是在使用 YOLOv5 进行目标检测时设置的。这个类让模型可以在多种软硬件环境中运行。该类的 init 构造函数接受模型的权重文件、设备类型等参数。根据权重文件的格式(.pt、.onnx、*.engine等),这个类将使用合适的后端进行目标检测,这里就看pytorch(pt)部分。

    if pt:  # PyTorch
        model = attempt_load(weights if isinstance(weights, list) else w, device=device, inplace=True, fuse=fuse)
        stride = max(int(model.stride.max()), 32)  # model stride
        names = model.module.names if hasattr(model, 'module') else model.names  # get class names
        model.half() if fp16 else model.float()
        self.model = model  # explicitly assign for to(), cpu(), cuda(), half()

如果待加载的模型是 PyTorch 模型,在明确设备已经准备就绪后,加载模型;然后根据 fp16 的取值,用半精度浮点数 (.half()) 或者单精度浮点数 (.float()) 表示模型。如果 fp16 为 True,将模型转为半精度模式以减少内存占用和提高计算速度。Float16需要的内存更小,速度可能更快(特别是在支持half-precision运算的GPU上),但精度较低。

    def forward(self, im, augment=False, visualize=False):
        # YOLOv5 MultiBackend inference
        b, ch, h, w = im.shape  # batch, channel, height, width
        if self.fp16 and im.dtype != torch.float16:
            im = im.half()  # to FP16
        if self.nhwc:
            im = im.permute(0, 2, 3, 1)  # torch BCHW to numpy BHWC shape(1,320,192,3)

        if self.pt:  # PyTorch
            y = self.model(im, augment=augment, visualize=visualize) if augment or visualize else self.model(im)

这个forward()函数,首先是获取图像维度,nhwc的作用是某些框架下的维度表示方式不同,调整为一致的维度表示。即batchheightwidth*channel的形式。之后就是根据特定的框架

  • class AutoShape(nn.Module)

这个类的主要目的是封装对输入的预处理、网络推断和后处理,以方便处理各种形式的输入数据。类级别的参数有

  • conf和iou分别是NMS的置信度阈值和IoU阈值。
  • agnostic是一个布尔值,表示NMS是否对类别进行评估。
  • multi_label表示NMS是否为每个box提供多标签。
  • classes是一个可选的列表,用于按类别进行过滤。
  • max_det是每张图片的最大检测数量。
  • amp决定是否使用自动混合精度(AMP)进行推理。
       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
            if isinstance(ims, torch.Tensor):  # torch
                with amp.autocast(autocast):
                    return self.model(ims.to(p.device).type_as(p), augment=augment)  # inference

这一部分代码的作用是输入检查和设置,如果在在GPU上就打开AMP,AMP一般用于在不牺牲计算精度的前提下提高计算速度。此外确保了数据和模型在同一设备上,且有相同的数据类型。
预处理部分就给出如下注释:

# 如果 ims 是一个 list 或 tuple,那么记录图像的数量为 n,同时保留图像列表
# 如果 ims 是一个单独的图像,那么 n 设置为 1,并创建包含这个图像的列表
n, ims = (len(ims), list(ims)) if isinstance(ims, (list, tuple)) else (1, [ims]) 

# 初始化空列表:shape0 用于存储原始图像的形状,shape1 用于存储推理后的图像形状,files 用于存储文件名
shape0, shape1, files = [], [], []  

for i, im in enumerate(ims):
    # 给图像文件命名
    f = f'image{i}' 
    # 如果输入的是一个文件名或者URL,则打开这个文件并将其转换为 numpy 数组,并记录文件名
    if isinstance(im, (str, Path)): 
        im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im), im
        im = np.asarray(exif_transpose(im))
    # 如果输入的是一个 PIL Image,则将其转换为 numpy 数组,并记录文件名
    elif isinstance(im, Image.Image):
        im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename', f) or f
    # 将文件路径转换为 .jpg 文件,获取文件名,并添加到 files 列表
    files.append(Path(f).with_suffix('.jpg').name) 
    # 如果图像的形状是 CHW (即通道在前),则进行转置得到 HWC 的形状
    if im.shape[0] < 5:  
        im = im.transpose((1, 2, 0)) 
    # 如果图像是三通道的,只保留前三个通道,否则将其转为 BGR 彩色图像
    im = im[..., :3] if im.ndim == 3 else cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)
    # 记录图像的高和宽
    s = im.shape[:2] 
    # 记录原始图像的形状
    shape0.append(s)  
    # 计算图像缩放的比例
    g = max(size) / max(s)  
    # 根据比例改变图像大小
    shape1.append([int(y * g) for y in s])
    # 如果图像的数据不是连续的,那么创建一个新的连续版本
    ims[i] = im if im.data.contiguous else np.ascontiguousarray(im) 
# 将图像的尺寸修改为 self.stride 的倍数
shape1 = [make_divisible(x, self.stride) for x in np.array(shape1).max(0)]
# 对图像进行缩放和填充操作
x = [letterbox(im, shape1, auto=False)[0] for im in ims]
# 将数据重排成为 BCHW 的形式
x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2))) 
# 将 numpy 数组转换为 torch tensor,放到设备 p 上,然后转换类型并将数值范围从 [0,255] 改为 [0,1]
x = torch.from_numpy(x).to(p.device).type_as(p) / 255 

后处理如下


# 使用一个上下文管理器来自动处理混合精度运行(如果 autocast 为 True,则会在此范围内转换数据类型以加速运算)
with amp.autocast(autocast):

    # 在下面这个范围内进行模型预测
    with dt[1]:
        # 对输入x进行模型前向传播,augment表示是否对数据进行扩增处理
        y = self.model(x, augment=augment)  # forward

    # 在下面这个范围内进行预测结果的后处理
    with dt[2]:
        # 进行非最大抑制(NMS),用于检测后处理,避免获得大量重复的候选框
        # self.dmb表示是否应用Deep Max Blur模糊方法,conf是检测信心阈值,iou是IoU阈值
        # classes表示需要检测的目标类别,agnostic表示是否对所有类别执行NMS
        # multi_label表示是否允许每个盒子有多个类标签,max_det是最大检测对象数
        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])

    # 返回一个Detections对象,包含了图像、预测结果、输入文件、时间距离、类别名称以及输入的shape
    return Detections(ims, y, files, dt, self.names, x.shape)

  • class Detections

主要是_run 方法:对推理结果进行处理,以及选择性地进行显示、保存、剪裁或渲染。它会遍历每个检测,并应用所需的操作。它也会打印每张图片的检测结果。

  • class Proto(nn.Module)

用于分割的一些代码

  • class Classify(nn.Module)

用于分类的一些代码

experimentalpy

experimental.py提供了YOLOv5对象检测框架的一部分以及一些实验模块,这个脚本主要是延伸和增强了YOLOv5的各种功能,以更好地实现目标检测任务

  • Sum: 这是一个神经网络模块,计算2个或多个层的加权和。如果weight参数为True,那么每层的输出将乘以一个权重,这些权重是训练过程中学习的参数。否则,简单地将所有输入叠加在一起。

  • MixConv2d: 这是的混合深度卷积的实现。其中,每一层有不同数量的卷积核(由参数k指定)。这些卷积核的分布可以是每组相同的通道数量(如果equal_ch=True)或者在所有组之间的卷积核的数量相等。

  • Ensemble: 这是一个模型的集合的类。Ensemble的forward方法,对输入执行所有模型的前向传播并将结果连接起来。这允许你将多个模型的输出作为一个整体处理,常见的技术叫做模型融合。

  • attempt_load: 这个函数加载一个或多个模型的权重,然后返回一个包含所有模型的集合。这里有一些兼容性的更新的设置使得加载新模型仍可以正常工作。最后,如果只加载一个模型,只返回那一个模型,否则返回模型集合。

tf.py

tf.py是用来测试YOLOv5模型的脚本是否可以的文件
这部分代码定义一些封装好的TensorFlow层,分别对应于YOLOv5的不同模块。比如BatchNormalization、Pad、Convolution、Depthwise Convolution、Depthwise Convolution Transpose、Focus Layer、Bottleneck Layer、Cross Convolution等。

之后利用定义的小模块构建YOLOv5模型并解析YOLOv5配置文件,此外还给出应用该模型的实例,作用就是在三个框架下分别创建模型并进行为空的推断,以确定模型在这三个框架中是否可以工作。

yolo.py

yolo.py定义YOLOv5模型的主体结构

  • class Detect(nn.Module)

Detect 类负责生成目标检测的输出,包括目标的坐标、类别置信度等信息。它将卷积层的输出转换为检测结果。循环遍历YOLO网络的每一层,每一个预测的特征图都会被转换为一个预测网格,并且对应的预测参数会被计算和调整来生成预测的边框和预测的类别得分。
推理和训练模式:Detect 类能够根据模型的当前模式(推理或训练)来处理不同的数据流。在推理模式下,它生成目标检测结果,而在训练模式下,它可能包括损失函数的计算。

  • class BaseModel(nn.Module):

这个类是 YOLOv5 模型的基础模型,提供了一些通用的功能。函数较多,直接给出注释

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

    """
    这是一个用于分析模型性能的内置函数。
    它计算每一层的浮点运算数量(FLOPs)、参数数量和运行时间,并根据是否为最后一层输出汇总信息
    dt[-1],即上个模块的时间;o,即上个模块的FLOPs计数转化为GFLOPs;m.np,即当前模块的参数数;m.type,即当前模块的类型。
    """
    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")

    """
    对模型执行融合操作的函数,
    它将神经网络中的卷积层(Conv2d)和批量标准化层(BatchNorm2d)融合为一个层,以提高效率
    具体来说,两层之间不需要存储中间输出了
    """
    def fuse(self):  # fuse model Conv2d() + BatchNorm2d() layers
        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)  # update conv
                delattr(m, 'bn')  # remove batchnorm
                m.forward = m.forward_fuse  # update forward
        self.info()
        return self

    def info(self, verbose=False, img_size=640):  # print model information
        model_info(self, verbose, img_size)

        def _apply(self, fn):
            # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
            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))
            return self
  • class DetectionModel(BaseModel):

这一部分函数比较多,直接写在注释里

class DetectionModel(BaseModel):
    # YOLOv5 detection model
    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None):  # model, input channels, number of classes
        super().__init__()
        if isinstance(cfg, dict):
            self.yaml = cfg  # model dict
        else:  # is *.yaml
            import yaml  # for torch hub
            self.yaml_file = Path(cfg).name
            with open(cfg, encoding='ascii', errors='ignore') as f:
                self.yaml = yaml.safe_load(f)  # model dict

        # Define model
        ch = self.yaml['ch'] = self.yaml.get('ch', ch)  # input channels
        if nc and nc != self.yaml['nc']:
            LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
            self.yaml['nc'] = nc  # override yaml value
        if anchors:
            LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}')
            self.yaml['anchors'] = round(anchors)  # override yaml value
        self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])  # model, savelist
        self.names = [str(i) for i in range(self.yaml['nc'])]  # default names
        self.inplace = self.yaml.get('inplace', True)

        # Build strides, anchors
        m = self.model[-1]  # Detect()
        if isinstance(m, (Detect, Segment)):
            s = 256  # 2x min stride
            m.inplace = self.inplace
            forward = lambda x: self.forward(x)[0] if isinstance(m, Segment) else self.forward(x)
            m.stride = torch.tensor([s / x.shape[-2] for x in forward(torch.zeros(1, ch, s, s))])  # forward
            check_anchor_order(m)
            m.anchors /= m.stride.view(-1, 1, 1)
            self.stride = m.stride
            self._initialize_biases()  # only run once

        # Init weights, biases
        initialize_weights(self)
        self.info()
        LOGGER.info('')

    def forward(self, x, augment=False, profile=False, visualize=False):
        if augment:
            return self._forward_augment(x)  # augmented inference, None
        return self._forward_once(x, profile, visualize)  # single-scale inference, train

    """
    在一个循环中,该代码使用 zip(s, f) 对每个缩放因子和翻转操作进行一次遍历
    执行翻转操作(如果存在),并更改图像的尺寸,将增强后的图像通过模型进行一次前向运算,调整在增强预处理之后边界框尺寸的变化
    将结果添加到输出列表
    """
    def _forward_augment(self, x):
        img_size = x.shape[-2:]  # height, width
        s = [1, 0.83, 0.67]  # scales
        f = [None, 3, None]  # flips (2-ud, 3-lr)
        y = []  # outputs
        for si, fi in zip(s, f):
            xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
            yi = self._forward_once(xi)[0]  # forward
            # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1])  # save
            yi = self._descale_pred(yi, fi, si, img_size)
            y.append(yi)
        y = self._clip_augmented(y)  # clip augmented tails
        return torch.cat(y, 1), None  # augmented inference, train

    """
    这段代码的主要目的是在增强推理后反向调整预测结果,也就是逆算增强操作
    它首先将预测的前四个元素(应该是对应目标的边界框坐标)除以scale进行反向缩放操作。
    然后,依据flips的值,反向进行翻转操作。如果flips为2,则它会将y坐标做上下翻转;如果flips为3,则会将x坐标做左右翻转
    """
    def _descale_pred(self, p, flips, scale, img_size):
        # de-scale predictions following augmented inference (inverse operation)
        if self.inplace:
            p[..., :4] /= scale  # de-scale
            if flips == 2:
                p[..., 1] = img_size[0] - p[..., 1]  # de-flip ud
            elif flips == 3:
                p[..., 0] = img_size[1] - p[..., 0]  # de-flip lr
        else:
            x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale  # de-scale
            if flips == 2:
                y = img_size[0] - y  # de-flip ud
            elif flips == 3:
                x = img_size[1] - x  # de-flip lr
            p = torch.cat((x, y, wh, p[..., 4:]), -1)
        return p

    """
    这个函数的目的是修剪增强推理的尾部。
    这是因为模型在训练阶段需要对图像进行数据增强以增强其泛化能力,而推理阶段是不需要这种数据增强的,所以需要反操作消除影响
    """
    def _clip_augmented(self, y):
        # Clip YOLOv5 augmented inference tails
        nl = self.model[-1].nl  # number of detection layers (P3-P5)
        g = sum(4 ** x for x in range(nl))  # grid points
        e = 1  # exclude layer count
        i = (y[0].shape[1] // g) * sum(4 ** x for x in range(e))  # indices
        y[0] = y[0][:, :-i]  # large
        i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e))  # indices
        y[-1] = y[-1][:, i:]  # small
        return y

    """
    它主要用于初始化模型中的偏置项(biases),这些初始化方式的目的是为了在训练初期更好的进行梯度反向传播,从而加速模型的收敛速度。
    在目标检测中,模型需要对不同类别和目标的置信度进行建模,以及对类别的概率进行估计
    这种偏置是通过先验知识以及一些默认的策略进行的
    """
    def _initialize_biases(self, cf=None):  # initialize biases into Detect(), cf is class frequency
        # https://arxiv.org/abs/1708.02002 section 3.3
        # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
        m = self.model[-1]  # Detect() module
        for mi, s in zip(m.m, m.stride):  # from
            b = mi.bias.view(m.na, -1)  # conv.bias(255) to (3,85)
            b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
            b.data[:, 5:5 + m.nc] += math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())  # cls
            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
  • def parse_model(d, ch):

这部分的作用是解析参数,使用给出的参数对层实例化


runs文件夹:

这个文件夹用来保存输出

文件结构

在这里插入图片描述


export.py

export.py是YOLOv5中的一个Python脚本,用于将训练好的模型导出为ONNX格式或TorchScript格式,以便在其他平台上进行推理。该脚本还可以将模型转换为TensorRT格式,以在NVIDIA GPU上进行加速推理。


benchmarks.py

benchmarks.py主要用于评估YOLOv5模型在不同导出格式下的性能表现,包括导出模型的大小、mAP(平均精度均值)、推断时间等指标。根据指定的命令行参数,可以选择性运行基准测试或导出格式测试。


hubconf.py

hubconf.py是用于加载YOLOv5模型并进行推理的Python脚本,如果创建新模型并且还要加载预训练权重 pretrained=True,则会下载预训练权重文件 (如果尚未下载),然后载入到创建的模型中
这个脚本主要两部分:

  • 模型加载的函数:
    这里包含了加载多种预置的YOLOv5模型(yolov5n, yolov5s, yolov5m, yolov5l , yolov5x, yolov5n6, yolov5s6, yolov5m6, yolov5l6, yolov5x6等)以及自定义模型(custom)的方法。

  • _create函数:
    这个函数用来创建或者加载一个YOLOv5模型。参数包含模型名字、是否加载预训练权重、输入通道数量、类别数目、是否应用YOLOv5的.autoshape()包装器、是否打印所有信息以及要使用的设备等。


detect.py

其中有一个with dt[0]的用法,是python上下文管理器的用法,属于是我没接触过的python用法,with 上下文管理器用于创建一个临时的运行环境,确保在进入和退出这个环境时资源得到正确地分配和释放。在Python中,with 语句通常与上下文管理器对象一起使用,这个对象必须实现 enterexit 方法。以下是一个简单的用法示例,使用 with 语句来管理文件的打开和关闭:

# 创建一个自定义的上下文管理器类
class MyFileManager:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_value, traceback):
        self.file.close()

# 使用上下文管理器来读取文件
with MyFileManager('example.txt', 'r') as file:
    contents = file.read()
    print(contents)

# 文件在退出上下文管理器后自动关闭,无需显式调用 file.close()

如上示例展现了上下文管理器释放资源的作用,但是不够形象,具体作用看下面的代码:

class MyContextManager:
    def __enter__(self):
        print("Enter context...")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Exit context...")
        
with MyContextManager() as x:
    print("In the context...")

输出

Enter context...
In the context...
Exit context...

train.py

  • 1
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值