【YOLOv8改进 - 注意力机制】ContextAggregation : 上下文聚合模块，捕捉局部和全局上下文，增强特征表示

YOLOv8目标检测创新改进与实战案例专栏

专栏目录： YOLOv8有效改进系列及项目实战目录 包含卷积，主干 注意力，检测头等创新机制 以及 各种目标检测分割项目实战案例

专栏链接: YOLOv8基础解析+创新改进+实战案例

介绍

摘要

卷积神经网络（CNNs）在计算机视觉中无处不在，具有众多高效的变体。最近，最初在自然语言处理（NLP）中引入的Transformers越来越多地被应用于计算机视觉领域。尽管早期采用者继续使用CNN骨干网络，但最新的网络是端到端的、无CNN的Transformer解决方案。一个最近令人惊讶的发现表明，基于简单多层感知机（MLP）的解决方案，即使没有传统的卷积或Transformer组件，也能生成有效的视觉表示。虽然CNNs、Transformers和MLP-Mixers可以被认为是完全不同的架构，但我们提供了一个统一的视角，表明它们实际上是神经网络堆栈中聚合空间上下文的一种更普遍的方法的特例。我们提出了CONTAINER（CONText AggregatIon NEtwoRk），这是一个多头上下文聚合的通用构建模块，能够像Transformers一样利用长程交互，同时仍然利用局部卷积操作的归纳偏差，从而实现更快的收敛速度，这通常在CNNs中可见。我们的CONTAINER架构在ImageNet上使用22M参数实现了82.7%的Top-1准确率，相比DeiT-Small提高了2.8个百分点，并且只需200个epoch就能收敛到79.9%的Top-1准确率。与基于Transformer的方法相比，后者在依赖更大输入图像分辨率的下游任务中不具备良好的扩展性，我们的高效网络CONTAINER-LIGHT可以应用于DETR、RetinaNet和Mask-RCNN等目标检测和实例分割网络，获得令人印象深刻的检测mAP分别为38.9、43.8、45.1，和掩码mAP为41.3，相比使用可比计算和参数规模的ResNet-50骨干分别提高了6.6、7.3、6.9和6.6个点。与DeiT在DINO框架上的自监督学习相比，我们的方法也取得了可喜的结果。代码已在GitHub发布。

基本原理

关键组成部分和机制

1. 输入表示：

   • 输入图像表示为 $X\in {\mathbb{R}}^{C\times H\times W}$，其中 $C$ 是通道数， $H\times W$ 是空间维度。该图像被展平为一系列标记 { X i ∈ R C ∣ i = 1 , … , N } \{X_i \in \mathbb{R}^C | i = 1, \ldots, N\} {Xi​∈RC∣i=1,…,N}，其中 $N=H\times W$。

2. 亲和矩阵：

   • 定义亲和矩阵 $A\in {\mathbb{R}}^{N\times N}$ 以表示上下文聚合的邻域。该矩阵决定了信息在特征空间中的传播方式。聚合函数可以表示为：
     $$Y=(AV)W_1+X,$$

     其中 $V\in {\mathbb{R}}^{N\times C}$ 是通过线性投影 $V=X{W}_2$ 从 $X$ 获得的变换版本。 $W_1$ 和 $W_2$ 是可学习参数。

3. 多头聚合：

   • 为了增加建模能力，使用多个亲和矩阵，创建几条上下文信息的路径。多头版本的聚合函数为：
     $$Y=\text{Concat}(A_1{V}_1,\dots ,A_M{V}_M)W_2+X,$$

     其中每个 $A_m$ （对 $m=1,\dots ,M$）是表示特征空间内不同关系的亲和矩阵，增强了相比单头版本的表示能力。

4. 不同架构中的上下文聚合：

   • Transformers：利用基于投影的query-key对之间相似性动态生成的亲和矩阵。自注意机制可以表示为：
     $$A_{sa}^m=\text{Softmax}\left(\frac{Q_m{K}_m^T}{\sqrt{C/M}}\right),$$
     其中 $Q_m$ 和 $K_m$ 是query和key投影。这种方法捕捉实例级信息，但计算密集。

   • 深度卷积：并行融合空间和通道信息，使用从卷积核导出的静态和稀疏亲和矩阵。这种方法高效且适合捕捉局部上下文。

   • MLP-Mixer：使用转置MLP操作进行上下文聚合，具有密集亲和矩阵：
     $$A_{mlp}=(W_{MLP}{)}^T,$$
     其中 $W_{MLP}$ 是可学习参数。

5. 统一框架：CONTAINER：

   • CONTAINER（上下文聚合网络）模块通过可学习参数结合静态和动态亲和矩阵，提供灵活而强大的上下文聚合方法。一般形式为：
     $$Y=(\alpha A(X)+\beta A)V{W}_2+X,$$
     其中 $A(X)$ 是动态亲和矩阵， $A$ 是静态亲和矩阵。 $\alpha$ 和 $\beta$ 是可学习参数。

核心代码

class ContextAggregation(nn.Module):


    def __init__(self, in_channels, reduction=1):
        super(ContextAggregation, self).__init__()
        self.in_channels = in_channels
        self.reduction = reduction
        self.inter_channels = max(in_channels // reduction, 1)

        conv_params = dict(kernel_size=1, act_cfg=None)

        self.a = ConvModule(in_channels, 1, **conv_params)
        self.k = ConvModule(in_channels, 1, **conv_params)
        self.v = ConvModule(in_channels, self.inter_channels, **conv_params)
        self.m = ConvModule(self.inter_channels, in_channels, **conv_params)
        self.init_weights()

    def init_weights(self):
        for m in (self.a, self.k, self.v):
            caffe2_xavier_init(m.conv)
        constant_init(self.m.conv, 0)

    def forward(self, x):

        n = x.size(0)
        c = self.inter_channels
        a = self.a(x).sigmoid()
        k = self.k(x).view(n, 1, -1, 1).softmax(2)
        v = self.v(x).view(n, 1, c, -1)
        y = torch.matmul(v, k).view(n, c, 1, 1)
        y = self.m(y) * a
        return x + y

task与yaml配置

详见：https://blog.csdn.net/shangyanaf/article/details/140664662