当前位置：首页 > news >正文

YOLOv8改进：SA注意力机制【注意力系列篇】（附详细的修改步骤，以及代码，与其他一些注意力机制相比，不仅准确度更高，而且模型更加轻量化。）

news 2024/9/23 6:36:47

如果实验环境尚未搭建成功，可以参考这篇文章 ->【YOLOv8超详细环境搭建以及模型训练（GPU版本）】

文章链接为：http://t.csdnimg.cn/8ZmAm

--------------------------------------------------------------------------------------------------------------------------

1.基本原理简介

《SA-Net: Shuffle Attention for Deep Convolutional Neural Networks》提出了一种新的注意力机制——Shuffle Attention（SA），旨在增强卷积神经网络（CNN）的特征表达能力，并且提高其在图像分类、目标检测等视觉任务中的表现。该方法通过创新的注意力机制设计，结合高效的通道和空间特征处理，为网络提供了更多的表征能力，同时保持了较低的计算开销。

Abstract：Attention mechanisms, which enable a neural network to accurately focus on all the relevant elements of the input, have become an essential component to improve the performance of deep neural networks. There are mainly two attention mechanisms widely used in computer vision studies, spatial attention and channel attention, which aim to capture the pixel-level pairwise relationship and channel dependency, respectively. Although fusing them together may achieve better performance than their individual implementations, it will inevitably increase the computational overhead. In this paper, we propose an efficient Shuffle Attention (SA) module to address this issue, which adopts Shuffle Units to combine two types of attention mechanisms effectively. Specifically, SA first groups channel dimensions into multiple sub-features before processing them in parallel. Then, for each sub-feature, SA utilizes a Shuffle Unit to depict feature dependencies in both spatial and channel dimensions. After that, all sub-features are aggregated and a “channel shuffle” operator is adopted to enable information communication between different sub-features. The proposed SA module is efficient yet effective, e.g., the parameters and computations of SA against the backbone ResNet50 are 300 vs. 25.56M and 2.76e-3 GFLOPs vs. 4.12 GFLOPs, respectively, and the performance boost is more than 1.34% in terms of Top-1 accuracy. Extensive experimental results on commonused benchmarks, including ImageNet-1k for classification, MS COCO for object detection, and instance segmentation, demonstrate that the proposed SA outperforms the current SOTA methods significantly by achieving higher accuracy while having lower model complexity.

摘要：注意机制使神经网络能够准确地聚焦于输入的所有相关元素，已成为提高深度神经网络性能的重要组成部分。在计算机视觉研究中，主要有两种广泛使用的注意机制，即空间注意力和通道注意力，它们分别旨在捕捉像素级的成对关系和通道依赖关系。虽然将它们合并可能比它们各自的实现方式效果更好，但这必然会增加计算开销。在本文中，我们提出了一种有效的Shuffle Attention（SA）模块来解决这个问题，该模块采用Shuffle Units有效地组合了两种类型的注意机制。具体而言，SA首先将通道维度分组为多个子特征，然后并行处理它们。然后，对于每个子特征，SA利用Shuffle Unit描述了空间和通道维度上的特征依赖关系。之后，所有子特征被聚合，并采用“通道洗牌”运算符以促进不同子特征之间的信息交流。所提出的SA模块既高效又有效，例如，与骨干ResNet50相比，SA的参数和计算分别为300对25.56M和2.76e-3 GFLOPs对4.12 GFLOPs，而性能提升在Top-1准确率方面超过了1.34%。对常用基准测试的广泛实验结果，包括ImageNet-1k用于分类，MS COCO用于目标检测和实例分割，表明所提出的SA在实现更高准确率的同时具有更低的模型复杂度方面显着优于当前的SOTA方法。

研究背景

卷积神经网络在计算机视觉领域表现出色，被广泛应用于图像分类、目标检测、语义分割等任务。然而，传统的卷积操作在处理不同区域的特征时，并没有区分出哪些特征更为重要，这导致了网络在对输入特征的提取和利用方面存在一定的局限性。为了弥补这一不足，注意力机制应运而生。注意力机制通过引入自适应的特征加权策略，使得网络能够动态调整对不同特征的关注程度，从而有效提高网络对输入数据的理解和表征能力。

现有的注意力机制，如SE-Net（Squeeze-and-Excitation Networks）和CBAM（Convolutional Block Attention Module），虽然在提升模型性能方面取得了显著的效果，但其设计通常伴随着较高的计算开销或复杂的模型结构。因此，如何在保证计算效率的同时，进一步增强注意力机制的特征表达能力，成为了研究的重点。

SA-Net的基本原理

SA-Net的核心创新在于提出了Shuffle Attention机制，该机制通过将打乱（shuffle）操作与通道和空间注意力机制相结合，能够更高效地提升CNN的特征提取能力。其工作原理可以分为以下几个部分：

1. 通道注意力机制

通道注意力机制是注意力模型的一个重要组成部分，它的作用是捕捉通道之间的依赖关系。在一个卷积层输出的特征图中，不同通道往往表示不同的特征或模式，而并非所有特征都对任务同等重要。通过计算每个通道的重要性权重，通道注意力机制能够动态调整这些特征的表达，从而增强对有用特征的关注，削弱无用或冗余的信息。

2. Shuffle操作

SA-Net的创新之处在于引入了“Shuffle”操作。该操作首先将特征通道分组，并在不同组之间打乱通道的排列顺序。打乱操作使得不同通道间的信息能够相互交互，这种信息交换可以进一步加强通道间的联系和信息整合。通过这种打乱和信息交换的方式，网络能够更有效地捕捉全局和局部的特征关系。

3. 空间注意力机制

除了通道注意力，SA-Net还结合了空间注意力机制。空间注意力的作用是通过计算特征图上不同位置的响应值，来确定图像中哪些位置对于当前任务更为重要。通过关注不同位置的特征，网络可以更好地理解图像中的关键区域和背景之间的关系，进一步增强模型的表征能力。

4. Shuffle Attention模块

SA-Net中的Shuffle Attention模块同时结合了通道注意力和空间注意力机制，形成了一个统一的注意力框架。具体而言，特征图首先通过通道打乱操作进行重组，然后应用通道注意力机制调整特征通道的权重，最后通过空间注意力机制对特征进行进一步优化。在这一过程中，网络能够自适应地调整特征的表达，从而在不同任务中获得更好的性能。

SA-Net的优势

与现有的注意力机制相比，SA-Net具有以下几个显著优势：

高效性：Shuffle Attention模块引入的打乱操作和通道与空间注意力的结合，虽然提升了网络的特征处理能力，但并没有显著增加模型的计算复杂度。相比于其他注意力机制，SA-Net在性能提升的同时，计算开销更低，非常适合在深度网络中应用。

轻量化设计：SA-Net的模块设计简单且高效，能够轻松地嵌入到现有的卷积神经网络中，无需复杂的结构调整。这种轻量化的设计保证了其在各种任务中的广泛应用。

性能提升：通过在多个数据集上进行实验，SA-Net在图像分类、目标检测等任务中均取得了显著的性能提升。尤其在深度网络中，Shuffle Attention能够有效增强网络的特征提取和表达能力，使其在处理复杂视觉任务时表现更加出色。

实验结果与应用

SA-Net在ImageNet等大型图像分类数据集上的实验结果显示，它能够在保持较低参数量和计算开销的情况下，实现显著的性能提升。此外，SA-Net在目标检测任务中的应用也取得了良好的效果，表明其注意力机制对网络的全局特征提取和局部细节处理均有很好的增强作用。

这一注意力机制的设计还可以推广到其他视觉任务，如图像分割、视频理解等。通过嵌入SA模块，深度学习模型能够在各类视觉任务中获得更强的表达能力。

结论

SA-Net通过引入Shuffle Attention机制，为卷积神经网络提供了一种高效且轻量化的注意力机制。其独特的打乱操作和通道、空间注意力的结合，既保证了特征之间的有效交互，又保持了计算开销的低廉。实验结果表明，SA-Net在多个视觉任务中都取得了显著的性能提升。该方法不仅为卷积神经网络的改进提供了新的思路，也为注意力机制在计算机视觉中的应用拓宽了研究方向。

根据下图，可以看出，与其他常见的注意力记机制，准确度更高，模型更加的轻量化

论文地址：https://arxiv.org/pdf/2102.00240.pdf

代码地址：https://github.com/wofmanaf/SA-Net

2.将SA模块加入到YOLOv8中

注意：SA 模块是一种即插即用的模块，可以在许多位置添加。不同的数据集可能需要在不同的位置添加 SA 模块，其效果也会有所不同。建议在不同位置进行多次实验以便比较效果。以下是我选择添加SA模块的位置，供大家参考，但不一定要完全按照这种方式添加。

2.1 方法一：在YOLOv8的主干网络（Backbone）中添加SA模块

2.1.1 步骤一

在【ultralytics/nn/modules】目录下新建一个sa.py的文件,添加SA模块代码。代码获取链接：https://mbd.pub/o/bread/mbd-ZpuWmpZq

最后在【ultralytics/nn/moduels/_init_.py】文件里面导入SA模块。

2.1.2 步骤二

打开【ultralytics/nn/tasks.py】文件，在文件的开头导入SA模块。

然后找到parse_model这个方法（可以通过搜索parse_model）。

最后在parse_model方法中找到for语句(for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]))，在for语句中再添加一个elif语句，将SA模块添加进去。

elif m in {SA}:args = [ch[f],  ch[f]]

添加的位置截图如下：

2.1.3 步骤三

修改模型。在【ultralytics/cfg/models/v8】目录下新建一个yolov8-sa.yaml网络结构配置文件。将SA模块添加到YOLOv8结构中（这里我将SA模块添加到主干网络的倒数第二层）。 代码获取链接：https://mbd.pub/o/bread/mbd-ZpuWmpZq

2.1.4 训练过程

为了方便训练，可以在项目目录下新建一个名为train.py的文件，这样以后只需执行train.py就可以开始训练模型了。 代码获取链接：https://mbd.pub/o/bread/mbd-ZpuWmpZq

（如果train.py文件已经存在，只需要修改其中模型和数据集的配置文件路径即可。我这里给出的train.py文件中只是列举了常用的超参数，如果需要修改其他的超参数，可以进入【ultralytics/cfg/default.yaml】文件修改即可。）

成功运行的网络结构截图如下：

开始训练模型的部分截图如下：

训练结果保存在【runs/detect】目录下。

2.2 方法二：在C2f模块的残差结构中添加SA模块

2.2.1 步骤一

在【ultralytics/nn/modules】目录下新建一个c2f_sa.py的文件。将SA模块添加到C2f模块的残差结构中。 代码获取链接：https://mbd.pub/o/bread/mbd-ZpuWmpZq

然后在【ultralytics/nn/moduels/_init_.py】文件里面导入Bottleneck_SA，C2f_SA模块。

2.2.2 步骤二

打开【ultralytics/nn/tasks.py】文件，在文件的开头导入C2f_SA模块。

找到parse_model这个方法（可以通过搜索parse_model）。

然后在parse_model方法中找到for语句(for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]))，在for语句中的两个位置添加C2f_SA模块。

2.2.3 步骤三

修改模型。在【ultralytics/cfg/models/v8】目录下新建一个yolov8-c2f-sa.yaml网络结构配置文件。（这里我将主干中的C2f模块替换为C2f_SA模块，除此之外，也可以将Neck的C2f模块替换为C2f_SA模块） 代码获取链接：https://mbd.pub/o/bread/mbd-ZpuWmpZq