YOLOv8改进 | 注意力机制 | 在主干网络中添加MHSA模块【原理+附完整代码】

💡💡💡本专栏所有程序均经过测试，可成功执行💡💡💡

多头自注意力机制（Multi-Head Self-Attention）是Transformer模型中的一个核心概念，它允许模型在处理序列数据时同时关注不同的位置和表示子空间。这种机制是“自注意力”（Self-Attention）的一种扩展，自注意力又称为内部注意力（Intra-Attention），是一种注意力机制，用于对序列中的每个位置进行加权，以便在编码每个位置时能够考虑到序列中的其他位置。文章在介绍主要的原理后，将手把手教学如何进行模块的代码添加和修改，并将修改后的完整代码放在文章的最后，方便大家一键运行，小白也可轻松上手实践。以帮助您更好地学习深度学习目标检测YOLO系列的挑战。

专栏地址：YOLOv8改进——更新各种有效涨点方法——点击即可跳转

1. 原理

2. 多头自注意力机制代码实现

2.1 将MHSA添加到YOLOv8代码中

2.2 更改init.py文件

2.3 添加yaml文件

2.4 在task.py中进行注册

2.5 执行程序

3. 完整代码分享

4. GFLOPs

5. 总结

1. 原理

多头自注意力机制（Multi-Head Self-Attention, MHSA）是深度学习中的一种机制，主要用于提升模型捕捉复杂关系和不同尺度特征的能力。它是自注意力机制的扩展和增强版本，广泛应用于Transformer模型中，如BERT和GPT等。以下是多头自注意力机制的主要原理：

自注意力机制

首先，了解自注意力机制（Self-Attention Mechanism）的基础原理非常重要。在自注意力机制中，输入序列的每个元素（通常是词或词向量）都会根据其与其他元素的相关性进行加权。具体步骤如下：

输入表示：假设输入序列为 $x_1, x_2, \ldots, x_n$ ，每个元素是一个d维的向量。
线性变换：将输入向量通过三个不同的线性变换，得到查询（Query）、键（Key）和值（Value）矩阵。 $Q = XW_Q, \quad K = XW_K, \quad V = XW_V$ 其中， $W_Q$ 、 $W_K$ 、 $W_V$ 是可训练的权重矩阵。
计算注意力得分：通过点积来计算查询和键之间的相似性（即注意力得分），并应用缩放处理。 $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$ 其中， $(d_k)$ 是键向量的维度， $\text{softmax}$ 用于将得分转化为概率分布。
加权求和：用注意力得分对值向量进行加权求和，得到最终的输出。

多头自注意力机制

多头自注意力机制通过引入多个注意力头，扩展了自注意力机制的表示能力。具体步骤如下：

多头注意力：将输入向量通过不同的线性变换，得到多个不同的查询、键和值。 $\text{head}_i = \text{Attention}(Q_i, K_i, V_i)$ 其中，每个头都有独立的权重矩阵 $(W{Q_i})$ 、 $(W{K_i})$ 、 $(W_{V_i})$ 。
并行计算：对每个头并行计算注意力，并得到多个不同的输出。 $\text{head}_i = \text{Attention}(XW_{Q_i}, XW_{K_i}, XW_{V_i})$
拼接和线性变换：将所有头的输出拼接在一起，并通过一个线性变换得到最终输出。 $\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \ldots, \text{head}_h)W_O$ 其中，是用于将拼接后的向量映射回原始维度的权重矩阵。

主要优点

捕捉多种特征：多头机制允许模型在不同的子空间中捕捉输入的多种特征和关系。
增强表示能力：通过多头注意力，模型可以同时关注输入序列的不同部分，提高表示的多样性和丰富性。
稳定训练：多头机制还可以缓解单头注意力可能出现的不稳定性问题。

总之，多头自注意力机制通过并行计算多个注意力头，有效增强了模型的表示能力，使其能够更好地捕捉序列数据中的复杂模式和关系。这一机制在自然语言处理和其他序列数据任务中表现出色，是Transformer模型成功的关键组件之一。

2. 多头自注意力机制代码实现

2.1 将MHSA添加到YOLOv8代码中

关键步骤一: 将下面代码粘贴到在/ultralytics/ultralytics/nn/modules/conv.py中，并在该文件的__all__中添加“MHSA”

class MHSA(nn.Module):
    def __init__(self, n_dims, width=14, height=14, heads=4, pos_emb=False):
        super(MHSA, self).__init__()

        self.heads = heads
        self.query = nn.Conv2d(n_dims, n_dims, kernel_size=1)
        self.key = nn.Conv2d(n_dims, n_dims, kernel_size=1)
        self.value = nn.Conv2d(n_dims, n_dims, kernel_size=1)
        self.pos = pos_emb
        if self.pos:
            self.rel_h_weight = nn.Parameter(torch.randn([1, heads, (n_dims) // heads, 1, int(height)]),
                                             requires_grad=True)
            self.rel_w_weight = nn.Parameter(torch.randn([1, heads, (n_dims) // heads, int(width), 1]),
                                             requires_grad=True)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        n_batch, C, width, height = x.size()
        q = self.query(x).view(n_batch, self.heads, C // self.heads, -1)
        k = self.key(x).view(n_batch, self.heads, C // self.heads, -1)
        v = self.value(x).view(n_batch, self.heads, C // self.heads, -1)
        content_content = torch.matmul(q.permute(0, 1, 3, 2), k)  # 1,C,h*w,h*w
        c1, c2, c3, c4 = content_content.size()
        if self.pos:
            content_position = (self.rel_h_weight + self.rel_w_weight).view(1, self.heads, C // self.heads, -1).permute(
                0, 1, 3, 2)  # 1,4,1024,64

            content_position = torch.matmul(content_position, q)  # ([1, 4, 1024, 256])
            content_position = content_position if (
                    content_content.shape == content_position.shape) else content_position[:, :, :c3, ]
            assert (content_content.shape == content_position.shape)
            energy = content_content + content_position
        else:
            energy = content_content
        attention = self.softmax(energy)
        out = torch.matmul(v, attention.permute(0, 1, 3, 2))  # 1,4,256,64
        out = out.view(n_batch, C, width, height)
        return out

多头自注意力机制在处理图片时，主要应用于视觉Transformer（Vision Transformer, ViT）等模型。这些模型将图片处理流程分为几个关键步骤，下面详细解释这些步骤：

1. 图像分块

输入图片：假设输入图像的尺寸为 $H \times W \times C$ ，其中H是高度，W是宽度，C是通道数（例如RGB图像的C=3）。

分块：将输入图像划分为多个固定大小的小块（patch）。假设每个小块的尺寸为 $P \times P \times C$ ，则图片会被划分成 $N = \frac{H \times W}{P \times P}$ 个小块。

2. 小块向量化

线性变换：将每个小块展平（flatten）为一个向量，然后通过线性变换将其映射到固定的维度，得到小块的嵌入表示。假设映射后的维度为D，则每个小块被表示为一个D维向量。

3. 加入位置编码

位置编码：由于Transformer不具备内建的位置感知能力，需加入位置编码（positional encoding）以保留小块在图像中的位置信息。位置编码与小块嵌入相加，形成带有位置信息的输入序列。

4. 多头自注意力机制

输入准备：将带有位置编码的小块嵌入向量作为Transformer的输入序列。

多头自注意力：

查询、键和值：将每个小块嵌入向量通过不同的线性变换，得到查询（Query）、键（Key）和值（Value）。
并行计算多个注意力头：每个注意力头独立地计算注意力得分，然后通过加权求和值向量，得到每个头的输出。
拼接和线性变换：将所有头的输出拼接在一起，并通过一个线性变换得到最终输出。

5. 叠加Transformer层

多层堆叠：重复上述的多头自注意力和前馈神经网络结构，通常会堆叠多个这样的Transformer层。每一层都会进一步处理和组合小块的表示，提取更高级别的特征。

6. 分类头

CLS标记：在输入序列中添加一个特殊的分类标记（CLS token），它的嵌入表示会被用于最终的分类任务。

全连接层：通过一个全连接层（或多层感知机）将CLS标记的最终表示映射到目标类别数。

优点

全局信息捕捉：多头自注意力机制可以捕捉图像中远距离像素之间的关系。
灵活性强：不依赖于卷积操作，可以更灵活地处理不同分辨率的图像。
表示能力强：多头机制允许模型在不同的子空间中捕捉多种特征。

2.2 更改init.py文件

关键步骤二：修改modules文件夹下的__init__.py文件，先导入函数

然后在下面的__all__中声明函数

2.3 添加yaml文件

关键步骤三：在/ultralytics/ultralytics/cfg/models/v8下面新建文件yolov8_MHSA.yaml文件，粘贴下面的内容

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

# Parameters
nc: 2  # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.33, 0.25, 1024]  # YOLOv8n summary: 225 layers,  3157200 parameters,  3157184 gradients,   8.9 GFLOPs
  s: [0.33, 0.50, 1024]  # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients,  28.8 GFLOPs
  m: [0.67, 0.75, 768]   # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients,  79.3 GFLOPs
  l: [1.00, 1.00, 512]   # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
  x: [1.00, 1.25, 512]   # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs

# YOLOv8.0n 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, 3, C2f, [128, True]]
  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
  - [-1, 6, C2f, [256, True]]
  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
  - [-1, 6, C2f, [512, True]]
  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
  - [-1, 3, C2f, [1024, True]]
  - [-1, 1, SPPF, [1024, 5]]  # 9
  - [-1, 1, MHSA, [14,14,4]]  #10

# YOLOv8.0n head
head:
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
  - [-1, 3, C2f, [512]]  # 13

  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
  - [-1, 3, C2f, [256]]  # 16 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 12], 1, Concat, [1]]  # cat head P4
  - [-1, 3, C2f, [512]]  # 19 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 9], 1, Concat, [1]]  # cat head P5
  - [-1, 3, C2f, [1024]]  # 22 (P5/32-large)

  - [[16, 19, 22], 1, Detect, [nc]]  # Detect(P3, P4, P5)

温馨提示：本文只是对yolov8基础上添加模块，如果要对yolov8n/l/m/x进行添加则只需要指定对应的depth_multiple 和 width_multiple。

# YOLOv8n
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.25  # layer channel multiple
 
# YOLOv8s
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
 
# YOLOv8l 
depth_multiple: 1.0  # model depth multiple
width_multiple: 1.0  # layer channel multiple
 
# YOLOv8m
depth_multiple: 0.67  # model depth multiple
width_multiple: 0.75  # layer channel multiple
 
# YOLOv8x
depth_multiple: 1.33  # model depth multiple
width_multiple: 1.25  # layer channel multiple

2.4 在task.py中进行注册

关键步骤四：在parse_model函数中进行注册，添加MHSA,

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

2.5 执行程序

关键步骤五: 在ultralytics文件中新建train.py，将model的参数路径设置为yolov8_MHSA.yaml的路径即可

建议大家写绝对路径，确保一定能找到

from ultralytics import YOLO
 
# Load a model
# model = YOLO('yolov8n.yaml')  # build a new model from YAML
# model = YOLO('yolov8n.pt')  # load a pretrained model (recommended for training)
 
model = YOLO(r'/projects/ultralytics/ultralytics/cfg/models/v8/yolov8_MHSA.yaml')  # build from YAML and transfer weights
 
# Train the model
model.train(device = [3], batch=16)

🚀运行程序，如果出现下面的内容则说明添加成功🚀

3. 完整代码分享

4. GFLOPs

关于GFLOPs的计算方式可以查看：百面算法工程师 | 卷积基础知识——Convolution

未改进的YOLOv8nGFLOPs

改进后的GFLOPs

5. 总结

自注意力机制是一种在处理序列数据时，通过计算序列中每个元素与其他元素之间的相关性来加权组合输入元素的方法。具体来说，自注意力机制首先对输入序列进行线性变换，生成查询（Query）、键（Key）和值（Value）向量。然后，通过计算查询和键的点积并进行缩放处理，再应用softmax函数，将相关性转换为概率分布，得到注意力得分。接着，用这些得分对值向量进行加权求和，得到每个元素的自注意力表示。这种机制允许模型在全局范围内捕捉输入序列中各个位置之间的依赖关系和特征。多头自注意力机制通过并行计算多个独立的自注意力，并将这些注意力头的输出拼接起来，再通过线性变换，进一步增强了模型捕捉多种特征和复杂关系的能力。这种机制在自然语言处理、计算机视觉等领域表现出色，如在Transformer模型中广泛应用，使得模型能够处理长距离依赖和捕捉全局上下文信息。