YOLOv9改进策略【注意力机制篇】| 引入MobileNetv4中的Mobile MQA，提高模型效率

本文主要是介绍YOLOv9改进策略【注意力机制篇】| 引入MobileNetv4中的Mobile MQA，提高模型效率，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

一、本文介绍

本文记录的是基于Mobile MQA模块的YOLOv9目标检测改进方法研究。MobileNetv4中的Mobile MQA模块是用于模型加速，减少内存访问的模块，相比其他全局的自注意力，其不仅加强了模型对全局信息的关注，同时也显著提高了模型效率。

文章目录

一、本文介绍
二、Mobile MQA注意力原理
三、Mobile MQA的实现代码
四、添加步骤
- 4.1 修改common.py
- - 4.1.1 创新模块⭐
- 4.2 修改yolo.py
五、yaml模型文件
- 5.1 模型改进⭐
六、成功运行结果

二、Mobile MQA注意力原理

在论文《MobileNetV4 - Universal Models for the Mobile Ecosystem》中，提出了Mobile MQA。

一、原理

基于MQA改进并结合不对称空间下采样：
- MQA（Multi-Query Attention）简化了传统的多头注意力机制，通过共享keys和values来减少内存访问需求。在移动混合模型中，当批量大小较小时，这种方式能有效提高运算强度。
- 借鉴MQA中对queries、keys和values的不对称计算方式，Mobile MQA引入了空间缩减注意力（SRA），对keys和values进行下采样，同时保持高分辨率的queries。这是因为在混合模型中，早期层的空间混合卷积滤波器使得空间上相邻的标记具有相关性。
- Mobile MQA的计算公式为：
  $Mobile\_MQA(X)= Concat(attention_1,...,attention_n)W^{O}$ ，
  其中 $attention_j = softmax(\frac{(XW^{Q_j})(SR(X)W^{K})^{T}}{\sqrt{d_k}})(SR(X)W^{V})$ ，这里SR可以是空间缩减操作（在设计中是一个步长为2的3x3深度卷积），也可以是恒等函数（当不进行空间缩减时）。

二、特点

针对加速器优化：专门为移动加速器进行了优化，考虑了移动加速器的计算和内存特性。
不对称空间下采样：通过对keys和values进行下采样，保持queries的高分辨率，在不损失太多精度的情况下，显著提高了效率。
操作简单高效：相比传统的注意力机制，Mobile MQA的设计更加简单，操作更加高效，更适合在移动设备上运行。

论文：http://arxiv.org/abs/2404.10518
源码：https://github.com/tensorflow/models/blob/master/official/vision/modeling/backbones/mobilenet.py

三、Mobile MQA的实现代码

Mobile MQA模块的实现代码如下：

class MultiQueryAttentionLayerWithDownSampling(nn.Module):def __init__(self, in_channels, num_heads, key_dim, value_dim, query_h_strides, query_w_strides, kv_strides, dw_kernel_size=3, dropout=0.0):"""Multi Query Attention with spatial downsampling.Referenced from here https://github.com/tensorflow/models/blob/master/official/vision/modeling/layers/nn_blocks.py3 parameters are introduced for the spatial downsampling:1. kv_strides: downsampling factor on Key and Values only.2. query_h_strides: vertical strides on Query only.3. query_w_strides: horizontal strides on Query only.This is an optimized version.1. Projections in Attention is explict written out as 1x1 Conv2D.2. Additional reshapes are introduced to bring a up to 3x speed up."""super(MultiQueryAttentionLayerWithDownSampling, self).__init__()self.num_heads = num_headsself.key_dim = key_dimself.value_dim = value_dimself.query_h_strides = query_h_stridesself.query_w_strides = query_w_stridesself.kv_strides = kv_stridesself.dw_kernel_size = dw_kernel_sizeself.dropout = dropoutself.head_dim = self.key_dim // num_headsif self.query_h_strides > 1 or self.query_w_strides > 1:self._query_downsampling_norm = nn.BatchNorm2d(in_channels)self._query_proj = conv2d(in_channels, self.num_heads * self.key_dim, 1, 1, norm=False, act=False)if self.kv_strides > 1:self._key_dw_conv = conv2d(in_channels, in_channels, dw_kernel_size, kv_strides, groups=in_channels,norm=True, act=False)self._value_dw_conv = conv2d(in_channels, in_channels, dw_kernel_size, kv_strides, groups=in_channels,norm=True, act=False)self._key_proj = conv2d(in_channels, key_dim, 1, 1, norm=False, act=False)self._value_proj = conv2d(in_channels, key_dim, 1, 1, norm=False, act=False)self._output_proj = conv2d(num_heads * key_dim, in_channels, 1, 1, norm=False, act=False)self.dropout = nn.Dropout(p=dropout)def forward(self, x):bs, seq_len, _, _ = x.size()# print(x.size())if self.query_h_strides > 1 or self.query_w_strides > 1:q = F.avg_pool2d(self.query_h_strides, self.query_w_strides)q = self._query_downsampling_norm(q)q = self._query_proj(q)else:q = self._query_proj(x)px = q.size(2)q = q.view(bs, self.num_heads, -1, self.key_dim)  # [batch_size, num_heads, seq_len, key_dim]if self.kv_strides > 1:k = self._key_dw_conv(x)k = self._key_proj(k)v = self._value_dw_conv(x)v = self._value_proj(v)else:k = self._key_proj(x)v = self._value_proj(x)k = k.view(bs, 1, self.key_dim, -1)   # [batch_size, 1, key_dim, seq_length]v = v.view(bs, 1, -1, self.key_dim)    # [batch_size, 1, seq_length, key_dim]# calculate attention score# print(q.shape, k.shape, v.shape)attn_score = torch.matmul(q, k) / (self.head_dim ** 0.5)attn_score = self.dropout(attn_score)attn_score = F.softmax(attn_score, dim=-1)# context = torch.einsum('bnhm,bmv->bnhv', attn_score, v)# print(attn_score.shape, v.shape)context = torch.matmul(attn_score, v)context = context.view(bs, self.num_heads * self.key_dim, px, px)output = self._output_proj(context)# print(output.shape)return output

参数	解释
in_channels	输入通道数
num_heads	自注意力头的数量
key_dim	键的维度
key_dim	值的维度
value_dim	仅用于查询的，在H方向上的步长
query_h_strides	仅用于查询的，在W方向上的步长
query_w_strides	仅对键和值进行下采样，1不进行下采样，2下采样
dw_kernel_size=3	深度可分离卷积的卷积核大小
dropout=0.0	随机丢失比例

四、添加步骤

4.1 修改common.py

此处需要修改的文件是models/common.py

common.py中定义了网络结构的通用模块，我们想要加入新的模块就只需要将模块代码放到这个文件内即可。

4.1.1 创新模块⭐

模块改进方法：基于Mobile MQA模块的RepNCSPELAN4。

改进方法是对YOLOv9中的RepNCSPELAN4模块进行改进。RepNCSPELAN4模块的创新思想是将CSP与ELAN相结合。CSP可以有效地分割梯度流，减少计算量的同时保持准确性。ELAN则通过灵活的层聚合方式，增强网络的学习能力。而MobileNetv4中的Mobile MQA模块则是用于模型加速，减少内存访问的模块，相比其他全局的自注意力，利用Mobile MQA替换RepNCSPELAN4模块中的卷积层后，不仅加强了模型对全局信息的关注，同时也显著提高了模型效率。

改进代码如下：参考代码

class MMQARepNCSPELAN4(nn.Module):# csp-elandef __init__(self, c1, c2, c3, c4, c5=1):  # ch_in, ch_out, number, shortcut, groups, expansionsuper().__init__()self.c = c3//2self.cv1 = Conv(c1, c3, 1, 1)self.cv2 = nn.Sequential(RepNCSP(c3//2, c4, c5), MultiQueryAttentionLayerWithDownSampling(c4, 8, 48, 48, 1, 1, 1))self.cv3 = nn.Sequential(RepNCSP(c4, c4, c5), MultiQueryAttentionLayerWithDownSampling(c4, 8, 48, 48, 1, 1, 1))self.cv4 = Conv(c3+(2*c4), c2, 1, 1)def forward(self, x):y = list(self.cv1(x).chunk(2, 1))y.extend((m(y[-1])) for m in [self.cv2, self.cv3])return self.cv4(torch.cat(y, 1))def forward_split(self, x):y = list(self.cv1(x).split((self.c, self.c), 1))y.extend(m(y[-1]) for m in [self.cv2, self.cv3])return self.cv4(torch.cat(y, 1))

在这里插入图片描述

注意❗：在4.2小节中的yolo.py文件中需要声明的模块名称为：MultiQueryAttentionLayerWithDownSampling和MMQARepNCSPELAN4。

4.2 修改yolo.py

此处需要修改的文件是models/yolo.py

yolo.py用于函数调用，我们只需要将common.py中定义的新的模块名添加到parse_model函数下即可。

MultiQueryAttentionLayerWithDownSampling模块以及MMQARepNCSPELAN4模块添加后如下：

在这里插入图片描述

五、yaml模型文件

5.1 模型改进⭐

在代码配置完成后，配置模型的YAML文件。

此处以models/detect/yolov9-c.yaml为例，在同目录下创建一个用于自己数据集训练的模型文件yolov9-c-MMQARepNCSPELAN4.yaml。

将yolov9-c.yaml中的内容复制到yolov9-c-MMQARepNCSPELAN4.yaml文件下，修改nc数量等于自己数据中目标的数量。

📌 模型的修改方法是将骨干网络中的所有RepNCSPELAN4模块替换成MMQARepNCSPELAN4模块，优化整体，提高效率。

结构如下：

# YOLOv9# parameters
nc: 1  # number of classes
depth_multiple: 1.0  # model depth multiple
width_multiple: 1.0  # layer channel multiple
#activation: nn.LeakyReLU(0.1)
#activation: nn.ReLU()# anchors
anchors: 3# YOLOv9 backbone
backbone:[[-1, 1, Silence, []],  # conv down[-1, 1, Conv, [64, 3, 2]],  # 1-P1/2# conv down[-1, 1, Conv, [128, 3, 2]],  # 2-P2/4# elan-1 block[-1, 1, MMQARepNCSPELAN4, [256, 128, 64, 1]],  # 3  修改此处# avg-conv down[-1, 1, ADown, [256]],  # 4-P3/8# elan-2 block[-1, 1, MMQARepNCSPELAN4, [512, 256, 128, 1]],  # 5  修改此处# avg-conv down[-1, 1, ADown, [512]],  # 6-P4/16# elan-2 block[-1, 1, MMQARepNCSPELAN4, [512, 512, 256, 1]],  # 7  修改此处# avg-conv down[-1, 1, ADown, [512]],  # 8-P5/32# elan-2 block[-1, 1, MMQARepNCSPELAN4, [512, 512, 256, 1]],  # 9  修改此处]# YOLOv9 head
head:[# elan-spp block[-1, 1, SPPELAN, [512, 256]],  # 10# up-concat merge[-1, 1, nn.Upsample, [None, 2, 'nearest']],[[-1, 7], 1, Concat, [1]],  # cat backbone P4# elan-2 block[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 13# up-concat merge[-1, 1, nn.Upsample, [None, 2, 'nearest']],[[-1, 5], 1, Concat, [1]],  # cat backbone P3# elan-2 block[-1, 1, RepNCSPELAN4, [256, 256, 128, 1]],  # 16 (P3/8-small)# avg-conv-down merge[-1, 1, ADown, [256]],[[-1, 13], 1, Concat, [1]],  # cat head P4# elan-2 block[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 19 (P4/16-medium)# avg-conv-down merge[-1, 1, ADown, [512]],[[-1, 10], 1, Concat, [1]],  # cat head P5# elan-2 block[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 22 (P5/32-large)# multi-level reversible auxiliary branch# routing[5, 1, CBLinear, [[256]]], # 23[7, 1, CBLinear, [[256, 512]]], # 24[9, 1, CBLinear, [[256, 512, 512]]], # 25# conv down[0, 1, Conv, [64, 3, 2]],  # 26-P1/2# conv down[-1, 1, Conv, [128, 3, 2]],  # 27-P2/4# elan-1 block[-1, 1, RepNCSPELAN4, [256, 128, 64, 1]],  # 28# avg-conv down fuse[-1, 1, ADown, [256]],  # 29-P3/8[[23, 24, 25, -1], 1, CBFuse, [[0, 0, 0]]], # 30  # elan-2 block[-1, 1, RepNCSPELAN4, [512, 256, 128, 1]],  # 31# avg-conv down fuse[-1, 1, ADown, [512]],  # 32-P4/16[[24, 25, -1], 1, CBFuse, [[1, 1]]], # 33 # elan-2 block[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 34# avg-conv down fuse[-1, 1, ADown, [512]],  # 35-P5/32[[25, -1], 1, CBFuse, [[2]]], # 36# elan-2 block[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 37# detection head# detect[[31, 34, 37, 16, 19, 22], 1, DualDDetect, [nc]],  # DualDDetect(A3, A4, A5, P3, P4, P5)]

六、成功运行结果

分别打印网络模型可以看到MMQARepNCSPELAN4已经加入到模型中，并可以进行训练了。

yolov9-c-MMQARepNCSPELAN4：

                 from  n    params  module                                  arguments                     0                -1  1         0  models.common.Silence                   []                            1                -1  1      1856  models.common.Conv                      [3, 64, 3, 2]                 2                -1  1     73984  models.common.Conv                      [64, 128, 3, 2]               3                -1  1    249472  models.common.MMQARepNCSPELAN4          [128, 256, 128, 64, 1]        4                -1  1    164352  models.common.ADown                     [256, 256]                    5                -1  1    773376  models.common.MMQARepNCSPELAN4          [256, 512, 256, 128, 1]       6                -1  1    656384  models.common.ADown                     [512, 512]                    7                -1  1   2119168  models.common.MMQARepNCSPELAN4          [512, 512, 512, 256, 1]       8                -1  1    656384  models.common.ADown                     [512, 512]                    9                -1  1   2119168  models.common.MMQARepNCSPELAN4          [512, 512, 512, 256, 1]       10                -1  1    656896  models.common.SPPELAN                   [512, 512, 256]               11                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          12           [-1, 7]  1         0  models.common.Concat                    [1]                           13                -1  1   3119616  models.common.RepNCSPELAN4              [1024, 512, 512, 256, 1]      14                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          15           [-1, 5]  1         0  models.common.Concat                    [1]                           16                -1  1    912640  models.common.RepNCSPELAN4              [1024, 256, 256, 128, 1]      17                -1  1    164352  models.common.ADown                     [256, 256]                    18          [-1, 13]  1         0  models.common.Concat                    [1]                           19                -1  1   2988544  models.common.RepNCSPELAN4              [768, 512, 512, 256, 1]       20                -1  1    656384  models.common.ADown                     [512, 512]                    21          [-1, 10]  1         0  models.common.Concat                    [1]                           22                -1  1   3119616  models.common.RepNCSPELAN4              [1024, 512, 512, 256, 1]      23                 5  1    131328  models.common.CBLinear                  [512, [256]]                  24                 7  1    393984  models.common.CBLinear                  [512, [256, 512]]             25                 9  1    656640  models.common.CBLinear                  [512, [256, 512, 512]]        26                 0  1      1856  models.common.Conv                      [3, 64, 3, 2]                 27                -1  1     73984  models.common.Conv                      [64, 128, 3, 2]               28                -1  1    212864  models.common.RepNCSPELAN4              [128, 256, 128, 64, 1]        29                -1  1    164352  models.common.ADown                     [256, 256]                    30  [23, 24, 25, -1]  1         0  models.common.CBFuse                    [[0, 0, 0]]                   31                -1  1    847616  models.common.RepNCSPELAN4              [256, 512, 256, 128, 1]       32                -1  1    656384  models.common.ADown                     [512, 512]                    33      [24, 25, -1]  1         0  models.common.CBFuse                    [[1, 1]]                      34                -1  1   2857472  models.common.RepNCSPELAN4              [512, 512, 512, 256, 1]       35                -1  1    656384  models.common.ADown                     [512, 512]                    36          [25, -1]  1         0  models.common.CBFuse                    [[2]]                         37                -1  1   2857472  models.common.RepNCSPELAN4              [512, 512, 512, 256, 1]       38[31, 34, 37, 16, 19, 22]  1  21542822  DualDDetect                             [1, [512, 512, 512, 256, 512, 512]]
yolov9-c-MMQARepNCSPELAN4 summary: 1018 layers, 49485350 parameters, 49485318 gradients, 236.8 GFLOPs

这篇关于YOLOv9改进策略【注意力机制篇】| 引入MobileNetv4中的Mobile MQA，提高模型效率的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！