YOLOv11改进 | 注意力篇 | YOLOv11引入MSDA多尺度空洞注意力

1. MSDA介绍

1.1 摘要：作为事实上的解决方案，鼓励使用普通视觉变换器（ViT）对任意图像块之间的远程依赖性进行建模，而全局参与感受野会导致二次计算成本。 Vision Transformers 的另一个分支利用了受 CNN 启发的局部注意力，它只模拟小邻域中斑块之间的相互作用。尽管这样的解决方案降低了计算成本，但它自然会受到较小的参与感受野的影响，这可能会限制性能。在这项工作中，我们探索有效的视觉变压器，以在计算复杂性和参与感受野的大小之间寻求更好的权衡。通过分析 ViT 中全局注意力的补丁交互，我们观察到浅层中的两个关键属性，即局部性和稀疏性，表明 ViT 浅层中全局依赖建模的冗余。因此，我们提出多尺度扩张注意力（MSDA）来模拟滑动窗口内的局部和稀疏补丁交互。通过金字塔架构，我们通过在低级阶段堆叠 MSDA 块和在高层阶段堆叠全局多头自注意力块来构建多尺度扩张变压器（DilateFormer）。我们的实验结果表明，我们的 DilateFormer 在各种视觉任务上实现了最先进的性能。在 ImageNet-1K 分类任务中，DilateFormer 实现了与现有最先进模型相当的性能，但 FLOP 数减少了 70%。我们的 DilateFormer-Base 在 ImageNet-1K 分类任务上实现了 85.6% 的 top-1 准确率，在 COCO 对象检测/实例分割任务上实现了 53.5% 的框 mAP/46.1% 的掩模 mAP，在 ADE20K 语义分割任务上实现了 51.1% 的 MS mIoU。

官方论文地址：https://arxiv.org/pdf/2302.01791

官方代码地址：https://github.com/JIAOJIAYUASD/dilateformer

1.2 简单介绍:

Multi-Scale Dilated Attention（MSDA）模块是DilateFormer中的核心组成部分，用于处理视觉任务中的长距离依赖关系。这一模块的设计灵感来源于自然语言处理领域，借鉴了序列模型中的成功经验。MSDA通过使用滑动窗口的方式来执行自注意力机制，与传统的全局注意力相比，它能够以较低的计算成本来模拟图像块之间的交互作用。

在DilateFormer架构中，MSDA模块被集成在浅层阶段，与深层阶段的多头自注意力（MHSA）模块形成对比。MSDA利用不同的膨胀率（dilation rates）为不同的头部设置不同的尺度，这使得该模块能够在保持特征图分辨率的同时增加感受野，并有效减少自注意力机制中的冗余。

具体来说，MSDA模块首先将输入的特征图划分为多个头部，然后在不同的头部中采用不同膨胀率的滑动窗口进行自注意力操作。这种多尺度策略允许模型同时捕获和融合多尺度的语义特征，从而提高了特征表示的丰富性和模型对细节的捕捉能力。

此外，MSDA模块的设计还考虑了局部性和稀疏性属性，即在浅层的注意力矩阵中，相关图像块通常在查询图像块的邻域内稀疏分布。这一点对于减少计算复杂度和内存使用具有重要意义，同时也有助于模型专注于与当前查询最相关的信息。

1.3 MSDA模块结构图

2. 核心代码

import torch
import torch.nn as nn
 
__all__ = ['MultiDilatelocalAttention']
 
 
class DilateAttention(nn.Module):
    "Implementation of Dilate-attention"
 
    def __init__(self, head_dim, qk_scale=None, attn_drop=0, kernel_size=3, dilation=1):
        super().__init__()
        self.head_dim = head_dim
        self.scale = qk_scale or head_dim ** -0.5
        self.kernel_size = kernel_size
        self.unfold = nn.Unfold(kernel_size, dilation, dilation * (kernel_size - 1) // 2, 1)
        self.attn_drop = nn.Dropout(attn_drop)
 
    def forward(self, q, k, v):
        # B, C//3, H, W
        B, d, H, W = q.shape
        q = q.reshape([B, d // self.head_dim, self.head_dim, 1, H * W]).permute(0, 1, 4, 3, 2)  # B,h,N,1,d
        k = self.unfold(k).reshape(
            [B, d // self.head_dim, self.head_dim, self.kernel_size * self.kernel_size, H * W]).permute(0, 1, 4, 2,
                                                                                                        3)  # B,h,N,d,k*k
        attn = (q @ k) * self.scale  # B,h,N,1,k*k
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        v = self.unfold(v).reshape(
            [B, d // self.head_dim, self.head_dim, self.kernel_size * self.kernel_size, H * W]).permute(0, 1, 4, 3,
                                                                                                        2)  # B,h,N,k*k,d
        x = (attn @ v).transpose(1, 2).reshape(B, H, W, d)
        return x
 
 
class MultiDilatelocalAttention(nn.Module):
    "Implementation of Dilate-attention"
 
    def __init__(self, dim, num_heads=8, qkv_bias=True, qk_scale=None,
                 attn_drop=0., proj_drop=0., kernel_size=3, dilation=[1, 2, 3, 4]):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.dilation = dilation
        self.kernel_size = kernel_size
        self.scale = qk_scale or head_dim ** -0.5
        self.num_dilation = len(dilation)
        assert num_heads % self.num_dilation == 0, f"num_heads{num_heads} must be the times of num_dilation{self.num_dilation}!!"
        self.qkv = nn.Conv2d(dim, dim * 3, 1, bias=qkv_bias)
        self.dilate_attention = nn.ModuleList(
            [DilateAttention(head_dim, qk_scale, attn_drop, kernel_size, dilation[i])
             for i in range(self.num_dilation)])
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
 
    def forward(self, x):
        B, C, H, W = x.shape
        # x = x.permute(0, 3, 1, 2)# B, C, H, W
        y = x.clone()
        qkv = self.qkv(x).reshape(B, 3, self.num_dilation, C // self.num_dilation, H, W).permute(2, 1, 0, 3, 4, 5)
        # num_dilation,3,B,C//num_dilation,H,W
        y1 = y.reshape(B, self.num_dilation, C // self.num_dilation, H, W).permute(1, 0, 3, 4, 2)
        # num_dilation, B, H, W, C//num_dilation
        for i in range(self.num_dilation):
            y1[i] = self.dilate_attention[i](qkv[i][0], qkv[i][1], qkv[i][2])  # B, H, W,C//num_dilation
        y2 = y1.permute(1, 2, 3, 0, 4).reshape(B, H, W, C)
        y3 = self.proj(y2)
        y4 = self.proj_drop(y3).permute(0, 3, 1, 2)
        return y4

3. YOLOv11中添加MSDA

3.1 在ultralytics/nn下新建Extramodule

3.2 在Extramodule里创建MSDA

在MSDA.py文件里添加给出的MSDA代码

添加完MSDA代码后，在ultralytics/nn/Extramodule/__init__.py文件中引用

3.3 在tasks.py里引用

在ultralytics/nn/tasks.py文件里引用Extramodule

在tasks.py找到parse_model（ctrl+f可以直接搜索parse_model位置）

添加如下代码：

        elif m in {MultiDilatelocalAttention}:
            c2 = ch[f]
            args = [c2, *args]

4. 新建一个yolo11MSDA.yaml文件

# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect

# Parameters
nc: 1 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs
  s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs
  m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs
  l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs
  x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs

# YOLO11n backbone
backbone:
  # [from, repeats, module, args]
  - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
  - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
  - [-1, 2, C3k2, [256, False, 0.25]]
  - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
  - [-1, 2, C3k2, [512, False, 0.25]]
  - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
  - [-1, 2, C3k2, [512, True]]
  - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
  - [-1, 2, C3k2, [1024, True]]
  - [-1, 1, SPPF, [1024, 5]] # 9
  - [-1, 2, C2PSA, [1024]] # 10

# YOLO11n head
head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 6], 1, Concat, [1]] # cat backbone P4
  - [-1, 2, C3k2, [512, False]] # 13
  - [-1, 1, MultiDilatelocalAttention, []]

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small)
  - [-1, 1, MultiDilatelocalAttention, []]

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 13], 1, Concat, [1]] # cat head P4
  - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium)
  - [-1, 1, MultiDilatelocalAttention, []]

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 10], 1, Concat, [1]] # cat head P5
  - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large)
  - [-1, 1, MultiDilatelocalAttention, []]

  - [[17, 21, 26], 1, Detect, [nc]] # Detect(P3, P4, P5)

大家根据自己的数据集实际情况，修改nc大小。

5. 模型训练

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO(r'D:\yolo\yolov11\ultralytics-main\datasets\yolo11MSDA.yaml')
    model.train(data=r'D:\yolo\yolov11\ultralytics-main\datasets\data.yaml',
                cache=False,
                imgsz=640,
                epochs=100,
                single_cls=False,  # 是否是单类别检测
                batch=4,
                close_mosaic=10,
                workers=0,
                device='0',
                optimizer='SGD',
                amp=True,
                project='runs/train',
                name='exp',
                )

模型结构打印，成功运行：