YOLO11改进|注意力机制篇|引入EMA注意力机制

如果能为勤奋颁奖

于 2024-10-09 09:52:17 发布

阅读量495

点赞数 8

分类专栏： YOLO11改进专栏文章标签： YOLO

本文链接：https://blog.csdn.net/A1983Z/article/details/142774172

版权

YOLO11改进专栏专栏收录该内容

25 篇文章 10 订阅

订阅专栏

在这里插入图片描述

一、【EMA】注意力机制

1.1【EMA】注意力介绍

在这里插入图片描述

下图是【EMA】的结构图，让我们简单分析一下运行过程和优势

处理过程分析：

Group 分组：
输入特征图的通道被分为多个组（Groups），以减少计算量。假设输入特征图形状为 𝐶×𝐻×𝑊，其中 𝐶是通道数，𝐻 和 𝑊是特征图的高和宽。通过分组，特征图被划分为 𝐺组，每组特征图的大小为 𝐶/𝐺×𝐻×𝑊。
X 方向和 Y 方向的全局池化：
X Avg Pool 和 Y Avg Pool：分别对特征图在 X（宽度）和 Y（高度）方向进行平均池化，得到两个大小为 𝐶/𝐺×𝐻×1和 𝐶/𝐺×1×𝑊的特征图。它们分别表示在 X 方向和 Y 方向上的全局特征。
Concat + Conv（1×1）：将两个特征图拼接后，通过 1×1 卷积进行特征融合，输出新的特征图，形状为 𝐶/𝐺×𝐻×𝑊，以进一步捕捉跨空间的信息。
Sigmoid 重新加权：
通过 Sigmoid 函数生成一个权重系数，接着对特征图进行 Re-weight（重新加权）。这一操作增强了具有重要特征的区域，同时抑制了不重要的区域。
Cross-spatial Learning（跨空间学习模块）：
GroupNorm：对每一组特征图进行归一化，保持每组特征在相同尺度下。
Avg Pool + Softmax + Matmul：这一部分是跨空间学习模块的核心部分，通过对每个特征组进行全局平均池化，结合 softmax 函数，生成跨空间注意力矩阵，然后通过矩阵乘法（Matmul）进行交互学习。结果输出为 1×𝐻×𝑊，捕捉了全局特征和位置间的依赖关系。
Sigmoid 加权：
经过跨空间学习后，生成的特征图通过 Sigmoid 进行非线性变换，产生加权系数，并再次对特征图进行加权，进一步强化重要信息。
输出：
最终加权后的特征图形状恢复为 𝐶×𝐻×𝑊，并输出给下游网络使用

优势分析：

跨空间特征捕捉：
通过 X 和 Y 方向的池化操作，该模块能够有效捕捉特征图中的空间信息，并结合跨空间的学习机制，使得网络能够捕获全局上下文中的依赖关系，从而增强对目标的识别和定位能力。
组卷积和归一化的高效性：
采用组卷积和 GroupNorm 机制，有效减少了计算量，同时保持较高的表达能力。分组处理的设计可以更好地分离通道特征和空间特征，提升网络的学习能力。
全局信息的整合：
模块通过全局池化和跨空间注意力机制，将图像的全局信息和局部特征结合在一起，使得模型能够关注到更广的上下文信息，从而提升对于大范围目标或复杂背景的感知能力。
逐步加权的精细调控：
通过多次 Sigmoid 加权操作，该模块在不同层次上对特征图进行精细的调控，能够逐层强化重要的特征区域，抑制不重要的信息，使得网络更加精确。
适应多尺度目标：
由于采用了多尺度的特征融合和跨空间学习，该模块能够在不同尺度下提取目标信息，非常适合处理多尺度物体或复杂场景中的目标检测任务。

1.2【EMA】核心代码

import torch
from torch import nn

# 定义一个 EMA 类，继承自 nn.Module。
class EMA(nn.Module):
    def __init__(self, channels, c2=None, factor=32):
        # 调用父类的构造函数。
        super(EMA, self).__init__()
        # 将通道数分成多个组，组的数量为 factor。
        self.groups = factor
        # 确保每组至少有一个通道。
        assert channels // self.groups > 0
        # 定义一个 softmax 层，用于计算权重。
        self.softmax = nn.Softmax(-1)
        # 自适应平均池化，将每个通道缩放到 (1,1) 的输出。
        self.agp = nn.AdaptiveAvgPool2d((1, 1))
        # 自适应平均池化，将高维度缩放到 1，宽维度保持不变。
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        # 自适应平均池化，将宽维度缩放到 1，高维度保持不变。
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))
        # 组归一化，每组的通道数为 channels // groups。
        self.gn = nn.GroupNorm(channels // self.groups, channels // self.groups)
        # 1x1 卷积层，用于通道的转换和维度缩减。
        self.conv1x1 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=1, stride=1, padding=0)
        # 3x3 卷积层，用于提取局部特征。
        self.conv3x3 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        # 获取输入 x 的形状 (batch_size, channels, height, width)。
        b, c, h, w = x.size()
        # 将输入 x 重新排列为 (batch_size * groups, channels // groups, height, width)。
        group_x = x.reshape(b * self.groups, -1, h, w)
        # 计算沿高度方向的池化，得到大小为 (batch_size * groups, channels // groups, height, 1) 的张量。
        x_h = self.pool_h(group_x)
        # 计算沿宽度方向的池化，得到大小为 (batch_size * groups, channels // groups, 1, width) 的张量，并进行转置。
        x_w = self.pool_w(group_x).permute(0, 1, 3, 2)
        # 将两个池化的张量连接在一起，并通过 1x1 卷积层，得到一个新的特征图。
        hw = self.conv1x1(torch.cat([x_h, x_w], dim=2))
        # 将特征图按原来的高度和宽度切分，分别得到 x_h 和 x_w。
        x_h, x_w = torch.split(hw, [h, w], dim=2)
        # 使用 sigmoid 激活函数和 x_h, x_w 调整 group_x 的特征值。
        x1 = self.gn(group_x * x_h.sigmoid() * x_w.permute(0, 1, 3, 2).sigmoid())
        # 使用 3x3 卷积层对 group_x 进行特征提取。
        x2 = self.conv3x3(group_x)
        # 计算 x1 的平均池化并通过 softmax，得到权重。
        x11 = self.softmax(self.agp(x1).reshape(b * self.groups, -1, 1).permute(0, 2, 1))
        # 将 x2 重新排列为 (batch_size * groups, channels // groups, height * width) 的形状。
        x12 = x2.reshape(b * self.groups, c // self.groups, -1)
        # 计算 x2 的平均池化并通过 softmax，得到权重。
        x21 = self.softmax(self.agp(x2).reshape(b * self.groups, -1, 1).permute(0, 2, 1))
        # 将 x1 重新排列为 (batch_size * groups, channels // groups, height * width) 的形状。
        x22 = x1.reshape(b * self.groups, c // self.groups, -1)
        # 计算 x11 和 x12, x21 和 x22 的矩阵乘法，并将结果 reshape 为 (batch_size * groups, 1, height, width)。
        weights = (torch.matmul(x11, x12) + torch.matmul(x21, x22)).reshape(b * self.groups, 1, h, w)
        # 使用权重调整 group_x 的特征，并 reshape 为原始的形状 (batch_size, channels, height, width)。
        return (group_x * weights.sigmoid()).reshape(b, c, h, w)

# 测试代码
if __name__ == '__main__':
    # 创建一个 EMA 模块实例，通道数为 64，使用 CUDA 加速。
    block = EMA(64).cuda()
    # 生成一个大小为 (1, 64, 64, 64) 的随机张量作为输入，并使用 CUDA 加速。
    input = torch.rand(1, 64, 64, 64).cuda()
    # 将输入张量传递给 EMA 模块，计算输出。
    output = block(input)
    # 打印输入和输出的形状，确保它们匹配。
    print(input.size(), output.size())

二、添加【EMA】注意力机制

2.1STEP1

首先找到ultralytics/nn文件路径下新建一个Add-module的python文件包【这里注意一定是python文件包，新建后会自动生成_init_.py】，如果已经跟着我的教程建立过一次了可以省略此步骤，随后新建一个EMA.py文件并将上文中提到的注意力机制的代码全部粘贴到此文件中，如下图所示在这里插入图片描述

2.2STEP2

在STEP1中新建的_init_.py文件中导入增加改进模块的代码包如下图所示在这里插入图片描述

2.3STEP3

找到ultralytics/nn文件夹中的task.py文件，在其中按照下图添加在这里插入图片描述

2.4STEP4

定位到ultralytics/nn文件夹中的task.py文件中的def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)函数添加如图代码,【如果不好定位可以直接ctrl+f搜索定位】

在这里插入图片描述

三、yaml文件与运行

3.1yaml文件

以下是添加【EMA】注意力机制在Backbone中的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: 80 # 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,1,EMA,[]]
  - [-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, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small)

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

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

  - [[17, 20, 23], 1, Detect, [nc]] # Detect(P3, P4, P5)