学习资源站

YOLOv11改进-添加注意力机制篇-添加MSDA多尺度空洞注意力(全新的YOLOv11改进策略,二次创新C2PSA机制)

一、本文介绍

本文给大家带来的改进机制是 MSDA(多尺度空洞注意力) 发表于中科院一区(算是国内计算机领域的最高期刊了) ,其全称是"DilateFormer: Multi-Scale Dilated Transformer for Visual Recognition"。MSDA的主要思想是通过线性投影得到 特征图 X的相应查询、键和值。然后,将特征图的通道分成n个不同的头部,并在不同的头部中以不同的扩张率执行多尺度SWDA来提高模型的处理效率和检测精度。 亲测在小目标检测和大尺度目标检测的数据集上都有大幅度的涨点效果(mAP直接涨了大概有0.06左右 ) 最后本文会手把手教你添加MSDA模块到网络结构中。



二、MSDA框架原理

论文地址: 官方论文地址点击即可跳转

代码地址: 官方代码地址点击即可跳转


在DilateFormer论文中, 多尺度扩张注意力(MSDA) 模块是为了利用自注意机制在不同尺度上的稀疏性。MSDA通过线性投影得到特征图X的相应查询、键和值。然后,将特征图的通道分成n个不同的头部,并在不同的头部中以不同的扩张率执行多尺度SWDA。具体来说,MSDA被公式化如下:对于每个头部i,进行SWDA操作,并且对所有的输出 {h_i} 进行连接后送入一个线性层进行特征聚合。通过为不同的头部设置不同的扩张率,MSDA能够在被关注的接受域内有效地聚合不同尺度的语义信息,并在不需要复杂操作和额外计算成本的情况下有效地减少自注意机制的冗余

MSDA(多尺度扩张注意力)模块的主要改进机制包括以下几点:

1. 多尺度特征提取: 通过不同头部的 自注意力机制 ,MSDA能够捕捉到多尺度的语义信息,这对于理解图像的不同抽象层次是非常重要的。

2. 稀疏性利用: MSDA利用了自注意力机制在不同尺度的稀疏性,降低了计算的冗余,同时保持了 性能

3. 头部通道分离: MSDA将特征图的通道分离成多个头部,每个头部处理不同的特征子集,这样可以 并行 处理,增强了模型的学习能力和效率。

4. 不同的扩张率: 通过在不同头部设置不同的扩张率,MSDA能够在各个头部关注不同尺度的特征,从而能更加全面地捕捉图像中的信息。

5. 特征聚合: MSDA的输出通过连接操作合并,并通过线性层进行特征聚合,这样可以整合各个头部学习到的信息,得到更丰富的特征表示。

这些改进使得MSDA在不增加额外计算成本的情况下,提高了自注意力机制的效率和效果。

这幅图展示了ViT-Small的第三个多头自注意力(Multi-Head Self-Attention, MHSA)块的注意力图的可视化。在每张图中,一个特定的查询块(红色框内的区域)被用来展示其它各个块对它的注意力程度。注意力图显示了具有高注意力得分的块在查询块周围稀疏分布,而其它块的注意力得分较低。

这张图展示了多尺度扩张注意力(MSDA)的工作原理。在MSDA中,特征图的通道首先被分割成不同的头部,然后每个头部内部使用不同的扩张率(dilation rates)r来执行自注意力操作。这些操作在围绕红色查询块的窗口内的彩色块之间进行。

图中的例子展示了三种不同的扩张率(r=1, 2, 3)( 这里需要注意的是咱们我的网络中需要改成四种的扩张率 ),它们分别对应不同的感受野大小(3x3, 5x5, 7x7)。每个头部的自注意力操作针对的是其对应的扩张率和感受野。这样,模型能够在不同的尺度上捕捉图像特征,这些特征随后被连接在一起,并送入一个线性层进行特征聚合。

这种设计允许模型在不同的尺度上理解图像,从而提高对图像内容的整体理解。通过这种方法,MSDA不仅可以捕捉局部细节,也能够感知到更广泛区域的上下文信息,增强了模型的表现力。


三、MSDA核心代码

下面的代码是MSDA的核心代码,我们将其复制导' ultralytics /nn/modules'目录下,在其中创建一个文件,我这里起名为Dilation然后粘贴进去,其余使用方式看章节四。

  1. import torch
  2. import torch.nn as nn
  3. __all__ = ['MultiDilatelocalAttention', 'C2PSA_MSDA']
  4. class DilateAttention(nn.Module):
  5. "Implementation of Dilate-attention"
  6. def __init__(self, head_dim, qk_scale=None, attn_drop=0, kernel_size=3, dilation=1):
  7. super().__init__()
  8. self.head_dim = head_dim
  9. self.scale = qk_scale or head_dim ** -0.5
  10. self.kernel_size = kernel_size
  11. self.unfold = nn.Unfold(kernel_size, dilation, dilation * (kernel_size - 1) // 2, 1)
  12. self.attn_drop = nn.Dropout(attn_drop)
  13. def forward(self, q, k, v):
  14. # B, C//3, H, W
  15. B, d, H, W = q.shape
  16. 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
  17. k = self.unfold(k).reshape(
  18. [B, d // self.head_dim, self.head_dim, self.kernel_size * self.kernel_size, H * W]).permute(0, 1, 4, 2,
  19. 3) # B,h,N,d,k*k
  20. attn = (q @ k) * self.scale # B,h,N,1,k*k
  21. attn = attn.softmax(dim=-1)
  22. attn = self.attn_drop(attn)
  23. v = self.unfold(v).reshape(
  24. [B, d // self.head_dim, self.head_dim, self.kernel_size * self.kernel_size, H * W]).permute(0, 1, 4, 3,
  25. 2) # B,h,N,k*k,d
  26. x = (attn @ v).transpose(1, 2).reshape(B, H, W, d)
  27. return x
  28. class MultiDilatelocalAttention(nn.Module):
  29. "Implementation of Dilate-attention"
  30. def __init__(self, dim, num_heads=8, qkv_bias=True, qk_scale=None,
  31. attn_drop=0., proj_drop=0., kernel_size=3, dilation=[1, 2, 3, 4]):
  32. super().__init__()
  33. self.dim = dim
  34. self.num_heads = num_heads
  35. head_dim = dim // num_heads
  36. self.dilation = dilation
  37. self.kernel_size = kernel_size
  38. self.scale = qk_scale or head_dim ** -0.5
  39. self.num_dilation = len(dilation)
  40. assert num_heads % self.num_dilation == 0, f"num_heads{num_heads} must be the times of num_dilation{self.num_dilation}!!"
  41. self.qkv = nn.Conv2d(dim, dim * 3, 1, bias=qkv_bias)
  42. self.dilate_attention = nn.ModuleList(
  43. [DilateAttention(head_dim, qk_scale, attn_drop, kernel_size, dilation[i])
  44. for i in range(self.num_dilation)])
  45. self.proj = nn.Linear(dim, dim)
  46. self.proj_drop = nn.Dropout(proj_drop)
  47. def forward(self, x):
  48. B, C, H, W = x.shape
  49. # x = x.permute(0, 3, 1, 2)# B, C, H, W
  50. y = x.clone()
  51. qkv = self.qkv(x).reshape(B, 3, self.num_dilation, C // self.num_dilation, H, W).permute(2, 1, 0, 3, 4, 5)
  52. # num_dilation,3,B,C//num_dilation,H,W
  53. y1 = y.reshape(B, self.num_dilation, C // self.num_dilation, H, W).permute(1, 0, 3, 4, 2)
  54. # num_dilation, B, H, W, C//num_dilation
  55. for i in range(self.num_dilation):
  56. y1[i] = self.dilate_attention[i](qkv[i][0], qkv[i][1], qkv[i][2]) # B, H, W,C//num_dilation
  57. y2 = y1.permute(1, 2, 3, 0, 4).reshape(B, H, W, C)
  58. y3 = self.proj(y2)
  59. y4 = self.proj_drop(y3).permute(0, 3, 1, 2)
  60. return y4
  61. def autopad(k, p=None, d=1): # kernel, padding, dilation
  62. """Pad to 'same' shape outputs."""
  63. if d > 1:
  64. k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size
  65. if p is None:
  66. p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
  67. return p
  68. class Conv(nn.Module):
  69. """Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)."""
  70. default_act = nn.SiLU() # default activation
  71. def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
  72. """Initialize Conv layer with given arguments including activation."""
  73. super().__init__()
  74. self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
  75. self.bn = nn.BatchNorm2d(c2)
  76. self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
  77. def forward(self, x):
  78. """Apply convolution, batch normalization and activation to input tensor."""
  79. return self.act(self.bn(self.conv(x)))
  80. def forward_fuse(self, x):
  81. """Perform transposed convolution of 2D data."""
  82. return self.act(self.conv(x))
  83. class PSABlock(nn.Module):
  84. """
  85. PSABlock class implementing a Position-Sensitive Attention block for neural networks.
  86. This class encapsulates the functionality for applying multi-head attention and feed-forward neural network layers
  87. with optional shortcut connections.
  88. Attributes:
  89. attn (Attention): Multi-head attention module.
  90. ffn (nn.Sequential): Feed-forward neural network module.
  91. add (bool): Flag indicating whether to add shortcut connections.
  92. Methods:
  93. forward: Performs a forward pass through the PSABlock, applying attention and feed-forward layers.
  94. Examples:
  95. Create a PSABlock and perform a forward pass
  96. >>> psablock = PSABlock(c=128, attn_ratio=0.5, num_heads=4, shortcut=True)
  97. >>> input_tensor = torch.randn(1, 128, 32, 32)
  98. >>> output_tensor = psablock(input_tensor)
  99. """
  100. def __init__(self, c, attn_ratio=0.5, num_heads=4, shortcut=True) -> None:
  101. """Initializes the PSABlock with attention and feed-forward layers for enhanced feature extraction."""
  102. super().__init__()
  103. self.attn = MultiDilatelocalAttention(c)
  104. self.ffn = nn.Sequential(Conv(c, c * 2, 1), Conv(c * 2, c, 1, act=False))
  105. self.add = shortcut
  106. def forward(self, x):
  107. """Executes a forward pass through PSABlock, applying attention and feed-forward layers to the input tensor."""
  108. x = x + self.attn(x) if self.add else self.attn(x)
  109. x = x + self.ffn(x) if self.add else self.ffn(x)
  110. return x
  111. class C2PSA_MSDA(nn.Module):
  112. """
  113. C2PSA module with attention mechanism for enhanced feature extraction and processing.
  114. This module implements a convolutional block with attention mechanisms to enhance feature extraction and processing
  115. capabilities. It includes a series of PSABlock modules for self-attention and feed-forward operations.
  116. Attributes:
  117. c (int): Number of hidden channels.
  118. cv1 (Conv): 1x1 convolution layer to reduce the number of input channels to 2*c.
  119. cv2 (Conv): 1x1 convolution layer to reduce the number of output channels to c.
  120. m (nn.Sequential): Sequential container of PSABlock modules for attention and feed-forward operations.
  121. Methods:
  122. forward: Performs a forward pass through the C2PSA module, applying attention and feed-forward operations.
  123. Notes:
  124. This module essentially is the same as PSA module, but refactored to allow stacking more PSABlock modules.
  125. Examples:
  126. >>> c2psa = C2PSA(c1=256, c2=256, n=3, e=0.5)
  127. >>> input_tensor = torch.randn(1, 256, 64, 64)
  128. >>> output_tensor = c2psa(input_tensor)
  129. """
  130. def __init__(self, c1, c2, n=1, e=0.5):
  131. """Initializes the C2PSA module with specified input/output channels, number of layers, and expansion ratio."""
  132. super().__init__()
  133. assert c1 == c2
  134. self.c = int(c1 * e)
  135. self.cv1 = Conv(c1, 2 * self.c, 1, 1)
  136. self.cv2 = Conv(2 * self.c, c1, 1)
  137. self.m = nn.Sequential(*(PSABlock(self.c, attn_ratio=0.5, num_heads=self.c // 64) for _ in range(n)))
  138. def forward(self, x):
  139. """Processes the input tensor 'x' through a series of PSA blocks and returns the transformed tensor."""
  140. a, b = self.cv1(x).split((self.c, self.c), dim=1)
  141. b = self.m(b)
  142. return self.cv2(torch.cat((a, b), 1))
  143. if __name__ == "__main__":
  144. # Generating Sample image
  145. image_size = (1, 64, 240, 240)
  146. image = torch.rand(*image_size)
  147. # Model
  148. mobilenet_v1 = C2PSA_MSDA(64, 64)
  149. out = mobilenet_v1(image)
  150. print(out.size())


四、手把手教你添加MSDA模块

4.1 修改一

第一还是建立文件,我们找到如下ultralytics/nn文件夹下建立一个目录名字呢就是'Addmodules'文件夹( 用群内的文件的话已经有了无需新建) !然后在其内部建立一个新的py文件将核心代码复制粘贴进去即可。


4.2 修改二

第二步我们在该目录下创建一个新的py文件名字为'__init__.py'( 用群内的文件的话已经有了无需新建) ,然后在其内部导入我们的检测头如下图所示。


4.3 修改三

第三步我门中到如下文件'ultralytics/nn/tasks.py'进行导入和注册我们的模块( 用群内的文件的话已经有了无需重新导入直接开始第四步即可)

从今天开始以后的教程就都统一成这个样子了,因为我默认大家用了我群内的文件来进行修改!!


4.4 修改四

按照我的添加在parse_model里添加即可,两个图片都是本文的机制大家按照图片进行添加即可!


到此就修改完成了,大家可以复制下面的yaml文件运行。


4.2 MSDA的yaml文件和训练截图


4.2.1 MSDA的yaml版本一(推荐)

此版本训练信息:YOLO11-C2PSA-MSDA summary: 324 layers, 2,609,435 parameters, 2,609,419 gradients, 6.5 GFLOPs

版本说明:优化C2PSA注意力机制

  1. # Ultralytics YOLO 🚀, AGPL-3.0 license
  2. # YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
  3. # Parameters
  4. nc: 80 # number of classes
  5. scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n'
  6. # [depth, width, max_channels]
  7. n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs
  8. s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs
  9. m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs
  10. l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs
  11. x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs
  12. # YOLO11n backbone
  13. backbone:
  14. # [from, repeats, module, args]
  15. - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
  16. - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
  17. - [-1, 2, C3k2, [256, False, 0.25]]
  18. - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
  19. - [-1, 2, C3k2, [512, False, 0.25]]
  20. - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
  21. - [-1, 2, C3k2, [512, True]]
  22. - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
  23. - [-1, 2, C3k2, [1024, True]]
  24. - [-1, 1, SPPF, [1024, 5]] # 9
  25. - [-1, 2, C2PSA_MSDA, [1024]] # 10
  26. # YOLO11n head
  27. head:
  28. - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  29. - [[-1, 6], 1, Concat, [1]] # cat backbone P4
  30. - [-1, 2, C3k2, [512, False]] # 13
  31. - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  32. - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  33. - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small)
  34. - [-1, 1, Conv, [256, 3, 2]]
  35. - [[-1, 13], 1, Concat, [1]] # cat head P4
  36. - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium)
  37. - [-1, 1, Conv, [512, 3, 2]]
  38. - [[-1, 10], 1, Concat, [1]] # cat head P5
  39. - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large)
  40. - [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)


4.2.2 MSDA的yaml版本二

添加的版本二具体那种适合你需要大家自己多做实验来尝试。

此版本训练信息:YOLO11-MSDA summary: 370 layers, 2,940,571 parameters, 2,940,555 gradients, 7.1 GFLOPs

  1. # Ultralytics YOLO 🚀, AGPL-3.0 license
  2. # YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
  3. # Parameters
  4. nc: 80 # number of classes
  5. scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n'
  6. # [depth, width, max_channels]
  7. n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs
  8. s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs
  9. m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs
  10. l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs
  11. x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs
  12. # YOLO11n backbone
  13. backbone:
  14. # [from, repeats, module, args]
  15. - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
  16. - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
  17. - [-1, 2, C3k2, [256, False, 0.25]]
  18. - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
  19. - [-1, 2, C3k2, [512, False, 0.25]]
  20. - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
  21. - [-1, 2, C3k2, [512, True]]
  22. - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
  23. - [-1, 2, C3k2, [1024, True]]
  24. - [-1, 1, SPPF, [1024, 5]] # 9
  25. - [-1, 2, C2PSA, [1024]] # 10
  26. # YOLO11n head
  27. head:
  28. - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  29. - [[-1, 6], 1, Concat, [1]] # cat backbone P4
  30. - [-1, 2, C3k2, [512, False]] # 13
  31. - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  32. - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  33. - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small)
  34. - [-1, 1, MultiDilatelocalAttention, []] # 17 (P3/8-small) 小目标检测层输出位置增加注意力机制
  35. - [-1, 1, Conv, [256, 3, 2]]
  36. - [[-1, 13], 1, Concat, [1]] # cat head P4
  37. - [-1, 2, C3k2, [512, False]] # 20 (P4/16-medium)
  38. - [-1, 1, MultiDilatelocalAttention, []] # 21 (P4/16-medium) 中目标检测层输出位置增加注意力机制
  39. - [-1, 1, Conv, [512, 3, 2]]
  40. - [[-1, 10], 1, Concat, [1]] # cat head P5
  41. - [-1, 2, C3k2, [1024, True]] # 24 (P5/32-large)
  42. - [-1, 1, MultiDilatelocalAttention, []] # 25 (P5/32-large) 大目标检测层输出位置增加注意力机制
  43. # 具体在那一层用注意力机制可以根据自己的数据集场景进行选择。
  44. # 如果你自己配置注意力位置注意from[17, 21, 25]位置要对应上对应的检测层!
  45. - [[17, 21, 25], 1, Detect, [nc]] # Detect(P3, P4, P5)

4.4 MSDA的训练过程截图


4.5 训练代码

  1. import warnings
  2. warnings.filterwarnings('ignore')
  3. from ultralytics import YOLO
  4. if __name__ == '__main__':
  5. model = YOLO('模型yaml文件地址')
  6. # 如何切换模型版本, 上面的ymal文件可以改为 yolov8s.yaml就是使用的v8s,
  7. # 类似某个改进的yaml文件名称为yolov8-XXX.yaml那么如果想使用其它版本就把上面的名称改为yolov8l-XXX.yaml即可(改的是上面YOLO中间的名字不是配置文件的)!
  8. # model.load('yolov8n.pt') # 是否加载预训练权重,科研不建议大家加载否则很难提升精度
  9. model.train(data=r"填写你数据集yaml文件地址",
  10. # 如果大家任务是其它的'ultralytics/cfg/default.yaml'找到这里修改task可以改成detect, segment, classify, pose
  11. cache=False,
  12. imgsz=640,
  13. epochs=150,
  14. single_cls=False, # 是否是单类别检测
  15. batch=4,
  16. close_mosaic=0,
  17. workers=0,
  18. device='0',
  19. optimizer='SGD', # using SGD
  20. # resume=True, # 这里是填写True
  21. amp=False, # 如果出现训练损失为Nan可以关闭amp
  22. project='runs/train',
  23. name='exp',
  24. )


五、本文总结

到此本文的正式分享内容就结束了,在这里给大家推荐我的YOLOv11改进有效涨点专栏,本专栏目前为新开的平均质量分98分,后期我会根据各种最新的前沿顶会进行论文复现,也会对一些老的改进机制进行补充, 目前本专栏免费阅读(暂时,大家尽早关注不迷路~) ,如果大家觉得本文帮助到你了,订阅本专栏,关注后续更多的更新~

​​