YoloV10改进策略：Neck改进|Efficient-RepGFPN，实时目标检测的王者Neck

AI浩

发布于 2024-10-22 13:24:02

2300

发布于 2024-10-22 13:24:02

文章被收录于专栏：AI智韵

摘要

在实时目标检测领域，Yolo系列模型一直以其高效和准确而著称。近日，我们成功将Efficient-RepGFPN模块引入YoloV10中，实现了显著的涨点效果。这一改进不仅进一步提升了YoloV10的检测精度，还保留了其原有的高效性能，为实时目标检测领域带来了新的突破。

完整链接：

https://blog.csdn.net/m0_47867638/article/details/142867421

Efficient-RepGFPN模块是DAMO-YOLO中提出的一种高效重参数化广义特征金字塔网络。它充分借鉴了广义特征金字塔网络（GFPN）的优点，并在此基础上进行了多项创新，以满足实时目标检测的设计要求。通过引入不同尺度特征图具有不同通道维度的设置、移除额外的上采样操作以及结合重参数化机制和高效层聚合网络（ELAN）的连接，Efficient-RepGFPN模块在保持低延迟的同时，实现了高精度的特征融合和信息传递。

将Efficient-RepGFPN模块引入YoloV10后，我们进行了大量的实验验证。结果表明，这一改进使得YoloV10的检测精度得到了显著提升，同时在推理速度上仍然保持了实时性。这一成果得益于Efficient-RepGFPN模块的高效特征融合能力和强大的信息处理能力，使得YoloV10能够更好地应对复杂场景下的目标检测任务。

相比其他目标检测模型，基于Efficient-RepGFPN的YoloV10具有诸多优点。首先，它在保持高效性能的同时，实现了更高的检测精度，为实时目标检测领域树立了新的标杆。其次，Efficient-RepGFPN模块的引入使得YoloV10能够更好地适应不同规模和不同复杂度的目标检测任务，具有更强的泛化能力。最后，基于Efficient-RepGFPN的YoloV10在训练和推理过程中都表现出了更好的稳定性和鲁棒性，为实际应用提供了有力保障。

论文翻译：《DAMO-YOLO：一份关于实时目标检测设计的报告》

在本报告中，我们提出了一种快速且准确的目标检测方法，名为DAMO-YOLO，其性能优于最先进的YOLO系列。DAMO-YOLO在YOLO的基础上进行了一些新技术扩展，包括神经架构搜索（NAS）、高效的重新参数化的广义泛函网络（RepGFPN）、带有对齐OTA标签分配的轻量级头部以及蒸馏增强。特别是，我们使用基于最大熵原理的MAE-NAS方法，在低延迟和高性能的约束下搜索我们的检测骨干，生成具有空间金字塔池化和焦点模块的ResNet-like/CSP-like结构。在设计颈部和头部时，我们遵循“大颈部，小头部”的规则。我们引入了带有加速女王融合的广义泛函网络来构建检测颈部，并通过高效的层聚合网络（ELAN）和重新参数化升级其CSPNet。然后我们研究了头部大小如何影响检测性能，并发现只有一个任务投影层的重型颈部会产生更好的结果。此外，提出了对齐OTA来解决标签分配中的错位问题，并引入了一个蒸馏方案以提高性能到更高的水平。基于这些新技术，我们构建了一套在不同规模上的一系列模型，以满足不同场景的需求。对于一般行业需求，我们提出了DAMO-YOLO-T/S/M/L。它们在COCO上的平均精度（mAP）分别为43.6/47.7/50.2/51.9，并且在T4 GPU上的延迟分别为2.78/3.83/5.62/7.95毫秒。此外，对于计算能力有限的边缘设备，我们还提出了DAMO-YOLO-Ns/Nm/Nl轻量级模型。它们在COCO上的平均精度（mAP）分别为32.3/38.2/40.5，并且在X86-CPU上的延迟分别为4.08/5.05/6.69毫秒。我们提出的通用和轻量级模型在其各自的应用场景中表现优于其他YOLO系列模型。代码可在https://github.com/tinyvision/damo-yolo获取。

1 引言

近年来，研究人员在目标检测方法上取得了巨大进展[11, 27, 24, 23, 1, 33, 9]。工业界追求具有实时限制的高性能目标检测方法，而研究人员则专注于设计具有高效网络架构[4, 16, 30, 29, 17]和先进训练阶段[21, 31, 17, 26, 10]的单阶段检测器[23, 24, 22, 26, 1]。特别是，YOLOv5/6/7[32, 18, 34]、YOLOX[9]和PP-YOLOE[38]在COCO数据集上实现了显著的精度-延迟权衡，使得YOLO系列目标检测方法在工业界得到了广泛应用。

尽管目标检测已经取得了巨大进步，但仍然可以引入新技术来进一步提升性能。首先，网络结构在目标检测中起着至关重要的作用。Darknet在YOLO历史的早期阶段占据主导地位[24, 25, 26, 1, 32, 9]。最近，一些工作开始为他们的检测器研究其他高效的网络，即YOLOv6[18]和YOLOv7[34]。然而，这些网络仍然是手动设计的。得益于神经架构搜索（Neural Architecture Search，简称）的发展，通过技术发现了许多对检测友好的网络结构[4, 16, 30]，这些结构在性能上大大优于以前的手动设计网络。因此，我们利用技术，并将[30]{ }^{1} 引入到我们的DAMO-YOLO中。是一种启发式且无需训练的神经架构搜索方法，它不依赖于超网，可用于构建不同规模的主干网络。它能够生成具有空间金字塔池化和焦点模块的类似 / 的结构。

其次，对于检测器而言，学习高层语义特征和低层空间特征之间的充分融合信息至关重要，这使得检测器的颈部成为整个框架的关键部分。其他工作[17, 10, 36, 31]也讨论了颈部的重要性。特征金字塔网络（Feature Pyramid Network，FPN）[10]已被证明在融合多尺度特征方面十分有效。广义特征金字塔网络（Generalized-FPN，GFPN）[17]通过一种新颖的女王融合（queen-fusion）策略对FPN进行了改进。在DAMO-YOLO中，我们设计了一种重参数化广义特征金字塔网络（Reparameterized Generalized-FPN，RepGFPN）。它基于GFPN，但涉及加速的女王融合、高效层聚合网络（Efficient Layer Aggregation Networks，ELAN）和重参数化技术。

为了平衡延迟和性能，我们进行了一系列实验来验证检测器颈部和头部的重要性，并发现“大颈部，小头部”的设计会带来更好的性能。因此，我们放弃了之前YOLO系列工作[24, 25, 26, 1, 32, 9, 38]中的检测器头部，只保留了一个任务投影层。节省的计算量被转移到了颈部部分。除了任务投影模块外，头部没有其他训练层，因此我们将其命名为ZeroHead。与我们的RepGFPN相结合，ZeroHead实现了最先进的性能，我们相信这将为其他研究人员提供一些启示。

此外，动态标签分配方法，如OTA[8]和TOOD[7]，广受好评，与静态标签分配[43]相比取得了显著改进。然而，这些工作中仍存在对齐问题。我们提出了一种更好的解决方案，称为AlignOTA，以平衡分类和回归的重要性，从而部分解决该问题。

最后，知识蒸馏（Knowledge Distillation，KD）已被证明在大型模型监督下能有效提升小型模型的性能。这项技术非常适合实时目标检测的设计。然而，在YOLO系列上应用KD有时并不能取得显著改进，因为超参数难以优化且特征携带过多噪声。在我们的DAMO-YOLO中，我们首次在所有尺寸的模型上，尤其是小型模型上，再次实现了蒸馏技术的卓越性能。

如图1所示，通过上述改进，我们提出了一系列通用且轻量级的模型，这些模型的性能远超当前最先进的技术水平。

综上所述，本文的贡献主要有三个方面：

本文提出了一种新的检测器，称为DAMO-YOLO，它基于YOLO但融入了更多新技术，包括MAE-NAS主干网络、RepGFPN颈部、ZeroHead、AlignedOTA和蒸馏增强。
在公共COCO数据集上，DAMO-YOLO在通用和轻量级类别中均优于最先进的检测器（如YOLO系列）。
在DAMO-YOLO（小/中/大）中呈现了一整套不同规模的模型，以支持不同的部署。代码和预训练模型可在https://github.com/tinyvision/damo-yolo获取，并支持ONNX和TensorRT。

在这里插入图片描述

2 DAMO-YOLO

本节中，我们将详细介绍DAMO-YOLO的每个模块，包括神经架构搜索（Neural Architecture Search，NAS）主干网络、高效的重参数化广义特征金字塔网络（Reparameterized Generalized-FPN，RepGFPN）颈部、ZeroHead、AlignedOTA标签分配和蒸馏增强。DAMO-YOLO的整体框架如图3所示。

2.1 MAE-NAS主干网络

以往，在实时场景中，设计者依赖浮点运算数（Flops）-平均精度（mAP）曲线作为评估模型性能的简单手段。然而，模型的Flops与延迟之间的关系并不总是一致。为了提升模型在工业部署中的实际性能，DAMO-YOLO在设计过程中优先考虑了延迟-mAP曲线。

基于这一设计原则，我们使用MAE-NAS[30]来获得不同延迟预算下的最优网络。MAE-NAS基于信息理论构建了一个替代代理，用于在不进行训练的情况下对初始化网络进行排序。因此，搜索过程仅需数小时，远低于训练成本。MAE-NAS提供了几个基本搜索块，如Mob-block、Res-block和CSP-block，如图2所示。Mob-block是MobileNetV3[14]块的一种变体，Res-block源自ResNet[12]，而CSP-block则源自CSPNet[35]。完整支持的块列表可在MAE-NAS存储库中找到。

我们发现，在不同规模的模型中应用不同类型的块，可以在实时推理中实现更好的权衡。在表1中列出了在DAMO-YOLO中，不同规模下CSP-Darknet与我们的MAE-NAS主干网络的性能比较。在表中，“MAE-Res”表示我们在MAE-NAS主干网络中应用了Res-block，“MAE-CSP”表示我们应用了CSP-block。此外，“S”（Small）和“M”（Medium）代表不同规模的主干网络。如表1所示，通过MAENAS技术获得的MAE-CSP在速度和准确性方面都优于手动设计的CSP-Darknet，证明了MAE-NAS技术的优越性。此外，从表1中我们可以观察到，在较小模型上使用Res-block可以在性能和速度之间实现比CSP-block更好的权衡，而在更大、更深的网络上使用CSP-block则可以显著优于Res-block。因此，在最终设置中，我们在“T”（Tiny）和“S”模型中使用“MAE-Res”，在“M”和“L”模型中使用“MAE-CSP”。

当处理只有有限计算能力或没有GPU可用的场景时，拥有满足严格计算和速度要求的模型至关重要。为了解决这个问题，我们设计了一系列使用Mob-Block的轻量级模型。Mob-block源自MobleNetV3[14]，可以显著降低模型计算量，并且对CPU设备友好。

2.2 高效重参数化广义特征金字塔网络（Efficient RepGFPN）

特征金字塔网络旨在聚合从主干网络中提取的不同分辨率的特征，这已被证明是目标检测中至关重要且有效的部分[10, 36, 31]。传统的FPN（特征金字塔网络）[10]引入了一条自上而下的路径来融合多尺度特征。考虑到单向信息流的局限性，PAFPN（路径聚合特征金字塔网络）[36]增加了一条额外的自下而上的路径聚合网络，但计算成本更高。BiFPN（双向特征金字塔网络）[31]移除了只有一条输入边的节点，并在同一级别的原始输入上添加了跳跃连接。在[17]中，提出了广义特征金字塔网络（Generalized-FPN，GFPN）作为颈部结构，并实现了最优性能，因为它可以充分交换高级语义信息和低级空间信息。在GFPN中，前一层和当前层的同级别特征都进行了多尺度融合。此外，的跨层连接提供了更有效的信息传输，可以扩展到更深层的网络。

当我们在现代YOLO系列模型中将改进后的PANet直接替换为GFPN时，我们获得了更高的精度。然而，基于GFPN的模型的延迟远高于基于改进后的PANet的模型。通过分析GFPN的结构，原因可以归结为以下几个方面：1）不同尺度的特征图共享相同的通道维度；2）queen-fusion（此处可能指某种特殊的融合方式，但并非标准术语）操作不能满足实时检测模型的要求；3）基于卷积的跨尺度特征融合不够高效。

基于GFPN，我们提出了一种新型的高效重参数化广义特征金字塔网络（Efficient-RepGFPN）以满足实时目标检测的设计要求，这主要基于以下见解：1）由于不同尺度特征图的浮点运算数（FLOPs）差异较大，在有限计算成本的约束下，很难控制每个尺度特征图共享的通道维度相同。因此，在我们的颈部结构的特征融合中，我们采用了不同尺度特征图具有不同通道维度的设置。表3比较了相同和不同通道的性能以及精度，以及颈部深度和宽度的权衡所带来的好处。我们可以看到，通过灵活控制不同尺度的通道数量，我们可以获得比所有尺度共享相同通道更高的精度。当深度等于3且宽度等于（96,192,384）时，获得了最佳性能。2）GFPN通过queenfusion增强了特征交互，但也带来了大量额外的上采样和下采样操作。这些上采样和下采样操作的好处进行了比较，结果如表4所示。我们可以看到，额外的上采样操作导致延迟增加了0.6毫秒，而精度提升仅为0.3 mAP，远低于额外下采样操作带来的好处。因此，在实时检测的约束下，我们移除了queenfusion中的额外上采样操作。3）在特征融合块中，我们首先用CSPNet替换了原始的基于3x3卷积的特征融合，并获得了4.2 mAP的提升。之后，我们通过结合重参数化机制和高效层聚合网络（ELAN）的连接[34]对CSPNet进行了升级。在没有带来额外巨大计算负担的情况下，我们实现了更高的精度。比较结果如表5所示。

2.3 ZeroHead与AlignOTA

在物体检测的最新进展中，解耦头被广泛使用[9, 18, 38]。采用解耦头后，这些模型实现了更高的平均精度（AP），但延迟也显著增加。为了平衡延迟和性能，我们进行了一系列实验来权衡颈部和头部的重要性，结果如表6所示。从实验中我们发现，“大颈部，小头部”会带来更好的性能。因此，我们舍弃了先前工作[9, 18, 38]中的解耦头，仅保留了一个任务投影层，即一个用于分类的线性层和一个用于回归的线性层。我们将我们的头部命名为ZeroHead，因为我们的头部没有其他训练层。ZeroHead可以最大程度地节省重型RepGFPN颈部的计算量。值得注意的是，ZeroHead本质上可以被视为一个耦合头，这与其他工作[9,32,18,38]中的解耦头有很大不同。

在头部之后的损失函数中，我们遵循GFocal[20]的方法，使用质量焦点损失（QFL）进行分类监督，并使用分布焦点损失（DFL）和GIOU损失进行回归监督。QFL鼓励学习分类和定位质量的联合表示。DFL通过将边界框的位置建模为一般分布，提供了更多信息和更精确的边界框估计。所提出的DAMO-YOLO的训练损失公式为：

除了头部和损失函数外，标签分配是检测器训练过程中的一个关键组件，它负责将分类和回归目标分配给预定义的锚框。最近，动态标签分配方法，如OTA[8]和TOOD[7]，广受好评，与静态标签分配方法[43]相比取得了显著改进。动态标签分配方法根据预测和真实值之间的分配成本来分配标签，例如OTA[8]。尽管损失函数中分类和回归的对齐已被广泛研究[7, 20]，但当前工作中很少提及标签分配中分类和回归的对齐。分类和回归的不对齐是静态分配方法中的常见问题[43]。尽管动态分配缓解了这个问题，但由于分类和回归损失的不平衡（例如交叉熵和IoU损失[40]），它仍然存在。为了解决这个问题，我们将焦点损失[22]引入分类成本中，并使用预测和真实框之间的IoU作为软标签，公式如下：

通过这种公式，我们能够为每个目标选择分类和回归对齐的样本。除了对齐的分配成本外，我们还遵循OTA[8]的方法，从全局角度形成对齐分配成本的解决方案。我们将我们的标签分配方法命名为AlignOTA。标签分配方法的比较如表7所示。我们可以看到，AlignOTA优于所有其他标签分配方法。

在这里插入图片描述

2.4 蒸馏增强

知识蒸馏（KD）[13]是进一步提升小型模型性能的有效方法。然而，在YOLO系列上应用KD有时并不能取得显著改进，因为超参数难以优化，且特征包含过多噪声。在DAMO-YOLO中，我们首先在各种尺寸的模型上，尤其是小型模型上，再次使蒸馏大放异彩。我们采用基于特征的蒸馏来传递暗知识，这可以在中间特征图[15]中蒸馏出识别和定位信息。我们进行了快速验证实验，为我们的DAMO-YOLO选择了一种合适的蒸馏方法。结果如表8所示。我们得出结论，CWD更适合我们的模型，而MGD由于复杂超参数导致其通用性不足，表现不如Mimicking。

我们提出的蒸馏策略分为两个阶段：1) 在第一阶段（284个epoch），我们的教师模型在强马赛克域上对学生模型进行蒸馏。面对具有挑战性的增强数据分布，在教师模型的指导下，学生可以更顺畅地进一步提取信息。2) 在第二阶段（16个epoch），学生在无马赛克域上进行微调。我们不在此阶段采用蒸馏的原因是，在这么短的时间内，教师的经验会损害学生在陌生域（即无马赛克域）中的表现。长期蒸馏会减弱这种损害，但成本高昂。因此，我们选择了折中方案，让学生模型独立起来。

在DAMO-YOLO中，蒸馏配备了两个高级增强功能：1) 对齐模块。一方面，它是一个线性投影层，用于将学生特征调整到与教师特征相同的分辨率（C, H, W）。另一方面，直接强迫学生逼近教师特征，与自适应模仿[37]相比，只能带来较小的增益。2) 通道动态温度。受PKD[2]的启发，我们对教师和学生特征进行归一化，以减弱实际值差异带来的影响。在减去均值后，每个通道的标准差将作为KL损失中的温度系数。

此外，我们还提出了两个关键观察结果，以更好地利用蒸馏。一是蒸馏与任务损失之间的平衡。如图4所示，当我们更专注于蒸馏（权重=10）时，学生模型的分类损失收敛速度变慢，产生了负面影响。因此，较小的损失权重（权重=0.5）对于在蒸馏和分类之间取得平衡是必要的。另一个是检测器的浅层头部。我们发现，适当减少头部深度有利于颈部特征蒸馏。原因是当最终输出与蒸馏特征图之间的差距更小时，蒸馏可以对决策产生更好的影响。

在实践中，我们使用DAMO-YOLO-S作为教师模型来蒸馏DAMO-YOLO-T，使用DAMO-YOLO-M作为教师模型来蒸馏DAMO-YOLO-S，使用DAMO-YOLO-L作为教师模型来蒸馏DAMO-YOLO-M，而DAMO-YOLO-L则通过自身进行蒸馏。

2.5 通用类DAMO-YOLO模型

在本节中，我们介绍了通用类DAMO-YOLO模型，该模型已在包括COCO、Objects365和OpenImage在内的多个大规模数据集上进行了训练。我们的模型融合了多项改进，以实现高精度和泛化能力。首先，我们通过创建一个统一的标签空间进行过滤，解决了不同数据集中重叠类别所带来的歧义问题，例如“mouse”在COCO/Objects365中指的是计算机鼠标，而在OpenImage中指的是啮齿动物。这将原始的945个类别（COCO中的80个，Objects365中的365个，OpenImage中的500个）减少到了701个[42]。其次，我们提出了一种多项式平滑后加权采样的方法，以解决数据集大小不平衡的问题，这种不平衡通常会导致批量采样偏向于较大的数据集。最后，我们实施了类别感知采样方法，以解决Objects365和OpenImage数据集中的长尾效应问题。这为包含数据点较少的类别的图像分配了更大的采样权重。

通过改进模型，我们在大规模数据集上训练了DAMO-YOLO-S，与COCO上的DAMO-YOLO-S基线相比，实现了的mAP提升。此外，该预训练模型可用于下游任务的微调。我们将其应用于VisDrone数据集，进行了300个训练周期，并实现了的mAP，超过了在COCO上预训练模型的性能。这些结果突出了大规模数据集训练的优势以及它为模型提供的鲁棒性。结果如表11所示。

3 实现细节

我们的模型使用SGD优化器进行了300个训练周期的训练。权重衰减和SGD动量分别为和0.9。初始学习率为0.4，批量大小为256，学习率根据余弦计划进行衰减。遵循YOLO系列[9, 32, 18, 34]模型，我们使用了指数移动平均（EMA）和分组权重衰减。为了增强数据多样性，Mosaic[1, 32]和Mixup[41]增强是常见的做法。然而，最近的研究进展[44, 3]表明，在目标检测中，适当设计的框级增强至关重要。受此启发，我们在图像级增强中应用了Mosaic和Mixup，并在图像级增强后采用SADA[3]的框级增强，以获得更稳健的增强效果。

4 与SOTA的比较

DAMO-YOLO发布了一系列通用模型和轻量级模型，以满足通用场景和资源受限的边缘场景的需求。

为了评估DAMO-YOLO通用模型与其他最先进模型的性能，我们使用T4-GPU和TensorRT-FP16推理引擎进行了比较分析。如表9所示，结果表明，DAMO-YOLO的通用模型在准确性和速度方面都优于竞争对手。此外，我们提出的蒸馏技术在准确性方面提供了显著提升。

为了评估轻量级模型的性能，我们使用Intel-8163 CPU和Openvino作为推理引擎进行了比较分析。如表2所示，DAMO-YOLO的轻量级模型在速度和准确性方面都取得了显著优势，远超竞争对手。

5 结论

本文提出了一种新的目标检测方法，称为DAMO-YOLO，其性能优于YOLO系列中的其他方法。其优势来源于新技术，包括MAE-NAS骨干网络、高效的RepGFPN颈部、ZeroHead、AlignedOTA标签分配和蒸馏增强。

代码

对代码做了适当的修改，代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from .conv import RepConv
from .block import SPPF

def conv_bn(in_channels, out_channels, kernel_size, stride, padding, groups=1):
    '''Basic cell for rep-style block, including conv and bn'''
    result = nn.Sequential()
    result.add_module(
        'conv',
        nn.Conv2d(in_channels=in_channels,
                  out_channels=out_channels,
                  kernel_size=kernel_size,
                  stride=stride,
                  padding=padding,
                  groups=groups,
                  bias=False))
    result.add_module('bn', nn.BatchNorm2d(num_features=out_channels))
    return result

class Swish(nn.Module):
    def __init__(self, inplace=True):
        super(Swish, self).__init__()
        self.inplace = inplace

    def forward(self, x):
        if self.inplace:
            x.mul_(F.sigmoid(x))
            return x
        else:
            return x * F.sigmoid(x)

def get_activation(name='silu', inplace=True):
    if name is None:
        return nn.Identity()

    if isinstance(name, str):
        if name == 'silu':
            module = nn.SiLU(inplace=inplace)
        elif name == 'relu':
            module = nn.ReLU(inplace=inplace)
        elif name == 'lrelu':
            module = nn.LeakyReLU(0.1, inplace=inplace)
        elif name == 'swish':
            module = Swish(inplace=inplace)
        elif name == 'hardsigmoid':
            module = nn.Hardsigmoid(inplace=inplace)
        elif name == 'identity':
            module = nn.Identity()
        else:
            raise AttributeError('Unsupported act type: {}'.format(name))
        return module

    elif isinstance(name, nn.Module):
        return name

    else:
        raise AttributeError('Unsupported act type: {}'.format(name))

def get_norm(name, out_channels, inplace=True):
    if name == 'bn':
        module = nn.BatchNorm2d(out_channels)
    else:
        raise NotImplementedError
    return module

class ConvBNAct(nn.Module):
    """A Conv2d -> Batchnorm -> silu/leaky relu block"""
    def __init__(
        self,
        in_channels,
        out_channels,
        ksize,
        stride=1,
        groups=1,
        bias=False,
        act='silu',
        norm='bn',
        reparam=False,
    ):
        super().__init__()
        # same padding
        pad = (ksize - 1) // 2
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=ksize,
            stride=stride,
            padding=pad,
            groups=groups,
            bias=bias,
        )
        if norm is not None:
            self.bn = get_norm(norm, out_channels, inplace=True)
        if act is not None:
            self.act = get_activation(act, inplace=True)
        self.with_norm = norm is not None
        self.with_act = act is not None

    def forward(self, x):
        x = self.conv(x)
        if self.with_norm:
            x = self.bn(x)
        if self.with_act:
            x = self.act(x)
        return x

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

class BasicBlock_3x3_Reverse(nn.Module):
    def __init__(self,
                 ch_in,
                 ch_hidden_ratio,
                 ch_out,
                 act='relu',
                 shortcut=True):
        super(BasicBlock_3x3_Reverse, self).__init__()
        assert ch_in == ch_out
        ch_hidden = int(ch_in * ch_hidden_ratio)
        self.conv1 = ConvBNAct(ch_hidden, ch_out, 3, stride=1, act=act)
        self.conv2 = RepConv(ch_in, ch_hidden, 3, s=1, act=act)
        self.shortcut = shortcut

    def forward(self, x):
        y = self.conv2(x)
        y = self.conv1(y)
        if self.shortcut:
            return x + y
        else:
            return y


class CSPStage(nn.Module):
    def __init__(self,
                 ch_in,
                 ch_out,
                 n,
                 block_fn='BasicBlock_3x3_Reverse',
                 ch_hidden_ratio=1.0,
                 act='silu',
                 spp=False):
        super(CSPStage, self).__init__()

        split_ratio = 2
        ch_first = int(ch_out // split_ratio)
        ch_mid = int(ch_out - ch_first)
        self.conv1 = ConvBNAct(ch_in, ch_first, 1, act=act)
        self.conv2 = ConvBNAct(ch_in, ch_mid, 1, act=act)
        self.convs = nn.Sequential()

        next_ch_in = ch_mid
        for i in range(n):
            if block_fn == 'BasicBlock_3x3_Reverse':
                self.convs.add_module(
                    str(i),
                    BasicBlock_3x3_Reverse(next_ch_in,
                                           ch_hidden_ratio,
                                           ch_mid,
                                           act=act,
                                           shortcut=True))
            else:
                raise NotImplementedError
            if i == (n - 1) // 2 and spp:
                self.convs.add_module(
                    'sppf', SPPF(ch_mid * 4, ch_mid, 5))
            next_ch_in = ch_mid
        self.conv3 = ConvBNAct(ch_mid * n + ch_first, ch_out, 1, act=act)

    def forward(self, x):
        y1 = self.conv1(x)
        y2 = self.conv2(x)

        mid_out = [y1]
        for conv in self.convs:
            y2 = conv(y2)
            mid_out.append(y2)
        y = torch.cat(mid_out, axis=1)
        y = self.conv3(y)
        return y

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2024-10-14，如有侵权请联系 cloudcommunity@tencent.com 删除

数据