【YOLOv8改进 - 注意力机制】S2Attention : 整合空间位移和分割注意力

YOLO大王

于 2024-07-16 17:28:07 发布

阅读量267

点赞数 7

文章标签： YOLO 目标跟踪人工智能目标检测 yolov8

原创文章，禁止任何形式转载！

本文链接：https://blog.csdn.net/shangyanaf/article/details/140472196

版权

YOLOv8目标检测创新改进与实战案例专栏

专栏目录： YOLOv8有效改进系列及项目实战目录包含卷积，主干注意力，检测头等创新机制以及各种目标检测分割项目实战案例

专栏链接: YOLOv8基础解析+创新改进+实战案例

介绍

摘要

最近，基于MLP的视觉骨干网络开始出现。与CNN和视觉Transformer相比，具有较少归纳偏差的MLP架构在图像识别中表现出竞争力。其中，采用直接空间移位操作的空间移位MLP（S2-MLP）比包括MLP-mixer和ResMLP在内的早期工作取得了更好的性能。最近，使用较小的补丁和金字塔结构，Vision Permutator（ViP）和Global Filter Network（GFNet）在性能上超过了S2-MLP。本文中，我们改进了S2-MLP视觉骨干网络。我们在通道维度上扩展特征图，并将扩展后的特征图分成若干部分，对分割部分进行不同的空间移位操作。同时，我们利用分割注意力操作融合这些分割部分。此外，类似于其他方法，我们采用了较小尺度的补丁并使用金字塔结构来提升图像识别的准确性。我们将改进后的空间移位MLP视觉骨干网络称为S2-MLPv2。使用55M参数，我们的中型模型S2-MLPv2-Medium在ImageNet-1K基准测试中使用224 × 224的图像，在没有自注意力和外部训练数据的情况下实现了83.6%的top-1准确率。

文章链接

论文地址：论文地址

**代码地址：**代码地址

基本原理

分割注意力模块概述

分割注意力（Split Attention）模块旨在融合来自不同操作的多个特征图。这个模块在ResNeSt网络中首次提出，并被Vision Permutator等模型采用，用于增强特征图的表现能力。

具体实现步骤

特征图的平均计算：
- 假设有 $k$ 个大小为 $\times c$ 的特征图 $[X_1, X_2, \cdots, X_K]$ ，其中 $(n$ ) 是patches的数量， $c$ 是通道数。
- 首先，计算这些特征图的平均值：
  $\sum_{k=1}^{K} 1X_k$
  其中， $\in \mathbb{R}^n$ 是一个全1的行向量。
通过MLP生成注意力权重：
- 计算得到的平均值 (a) 会经过一系列的MLP层生成注意力权重：
  $\hat{a} = \sigma(aW_1)W_2$
  其中， $\sigma$ 是GELU激活函数， $KaTeX parse error: Can't use function '\)' in math mode at position 38: …\times \bar{c}}\̲)̲ 和 \(W_2 \in \m…$ 是MLP层的权重矩阵， $\hat{a} \in \mathbb{R}^{Kc}$ 是输出。
计算注意力矩阵：
- 输出 (\hat{a}) 会被重塑为矩阵 (\hat{A} \in \mathbb{R}^{K \times c})，然后通过softmax函数计算得到注意力矩阵：
  $\bar{A} = \text{softmax}(\hat{A})$
  其中，(\bar{A} \in \mathbb{R}^{K \times c})。
生成注意力特征图：
- 通过将每个特征图与注意力矩阵进行逐元素乘法，生成加权后的特征图 (\hat{X})：
  $\hat{X}[i, :] = \sum_{k=1}^{K} X_k[i, :] \odot \bar{A}[k, :]$
  其中，(\odot) 表示逐元素乘法。

论文中的实现

在论文的S2-MLPv2模块中，这一分割注意力机制被用来融合通过不同空间移位操作后的特征图。具体实现如下：

扩展输入特征图 (X) 的通道数，从 (c) 扩展到 (3c)：
$\hat{X} = \text{MLP}_1(X) \in \mathbb{R}^{w \times h \times 3c}$
将扩展后的特征图均等分成三部分：
$X^1 = \hat{X}[:, :, 1:c], \quad X^2 = \hat{X}[:, :, c+1:2c], \quad X^3 = \hat{X}[:, :, 2c+1:3c]$
分别对 (X^1) 和 (X^2) 进行空间移位操作，而 (X^3) 保持不变：
通过分割注意力模块融合这些特征图生成输出特征图：
$\hat{X} = \text{SA}(\{X^k\}_{k=1}^3)$
通过另一个MLP层生成最终输出：
$\bar{X} = \text{MLP}_2(\hat{X})$

核心代码

import numpy as np
import torch
from torch import nn
from torch.nn import init


def spatial_shift1(x):
    # 实现第一种空间位移，位移图像的四分之一块
    b, w, h, c = x.size()
    # 以下四行代码分别向左、向右、向上、向下移动图像的四分之一块
    x[:, 1:, :, :c // 4] = x[:, :w - 1, :, :c // 4]
    x[:, :w - 1, :, c // 4:c // 2] = x[:, 1:, :, c // 4:c // 2]
    x[:, :, 1:, c // 2:c * 3 // 4] = x[:, :, :h - 1, c // 2:c * 3 // 4]
    x[:, :, :h - 1, 3 * c // 4:] = x[:, :, 1:, 3 * c // 4:]
    return x


def spatial_shift2(x):
    # 实现第二种空间位移，逻辑与spatial_shift1相似，但位移方向不同
    b, w, h, c = x.size()
    # 对图像的四分之一块进行空间位移
    x[:, :, 1:, :c // 4] = x[:, :, :h - 1, :c // 4]
    x[:, :, :h - 1, c // 4:c // 2] = x[:, :, 1:, c // 4:c // 2]
    x[:, 1:, :, c // 2:c * 3 // 4] = x[:, :w - 1, :, c // 2:c * 3 // 4]
    x[:, :w - 1, :, 3 * c // 4:] = x[:, 1:, :, 3 * c // 4:]
    return x


class SplitAttention(nn.Module):
    # 定义分割注意力模块，使用MLP层进行特征转换和注意力权重计算
    def __init__(self, channel=512, k=3):
        super().__init__()
        self.channel = channel
        self.k = k  # 分割的块数
        # 定义MLP层和激活函数
        self.mlp1 = nn.Linear(channel, channel, bias=False)
        self.gelu = nn.GELU()
        self.mlp2 = nn.Linear(channel, channel * k, bias=False)
        self.softmax = nn.Softmax(1)

    def forward(self, x_all):
        # 计算分割注意力，并应用于输入特征
        b, k, h, w, c = x_all.shape
        x_all = x_all.reshape(b, k, -1, c)  # 重塑维度
        a = torch.sum(torch.sum(x_all, 1), 1)  # 聚合特征
        hat_a = self.mlp2(self.gelu(self.mlp1(a)))  # 通过MLP计算注意力权重
        hat_a = hat_a.reshape(b, self.k, c)  # 调整形状
        bar_a = self.softmax(hat_a)  # 应用softmax获取注意力分布
        attention = bar_a.unsqueeze(-2)  # 增加维度
        out = attention * x_all  # 将注意力权重应用于特征
        out = torch.sum(out, 1).reshape(b, h, w, c)  # 聚合并调整形状
        return out


class S2Attention(nn.Module):
    # S2注意力模块，整合空间位移和分割注意力
    def __init__(self, channels=512):
        super().__init__()
        # 定义MLP层
        self.mlp1 = nn.Linear(channels, channels * 3)
        self.mlp2 = nn.Linear(channels, channels)
        self.split_attention = SplitAttention()

    def forward(self, x):
        b, c, w, h = x.size()
        x = x.permute(0, 2, 3, 1)  # 调整维度顺序
        x = self.mlp1(x)  # 通过MLP层扩展特征
        x1 = spatial_shift1(x[:, :, :, :c])  # 应用第一种空间位移
        x2 = spatial_shift2(x[:, :, :, c:c * 2])  # 应用第二种空间位移
        x3 = x[:, :, :, c * 2:]  # 保留原始特征的一部分
        x_all = torch.stack([x1, x2, x3], 1)  # 堆叠特征
        a = self.split_attention(x_all)  # 应用分割注意力
        x = self.mlp2(a)  # 通过另一个MLP层缩减特征维度
        x = x.permute(0, 3, 1, 2)  # 调整维度顺序回原始
        return x


# 示例代码
if __name__ == '__main__':
    input = torch.randn(50, 512, 7, 7)  # 创建输入张量
    s2att = S2Attention(channels=512)  # 实例化S2注意力模块
    output = s2att(input)  # 通过S2注意力模块处理输入
    print(output.shape)  # 打印输出张量的形状

下载YoloV8代码

直接下载

GitHub地址

Git Clone

git clone https://github.com/ultralytics/ultralytics

安装环境

进入代码根目录并安装依赖。

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

在最新版本中，官方已经废弃了requirements.txt文件，转而将所有必要的代码和依赖整合进了ultralytics包中。因此，用户只需安装这个单一的ultralytics库，就能获得所需的全部功能和环境依赖。

pip install ultralytics

引入代码

在根目录下的ultralytics/nn/目录，新建一个 attention目录，然后新建一个以 S2Attention为文件名的py文件，把代码拷贝进去。

import numpy as np
import torch
from torch import nn
from torch.nn import init


def spatial_shift1(x):
    b, w, h, c = x.size()
    x[:, 1:, :, :c // 4] = x[:, :w - 1, :, :c // 4]
    x[:, :w - 1, :, c // 4:c // 2] = x[:, 1:, :, c // 4:c // 2]
    x[:, :, 1:, c // 2:c * 3 // 4] = x[:, :, :h - 1, c // 2:c * 3 // 4]
    x[:, :, :h - 1, 3 * c // 4:] = x[:, :, 1:, 3 * c // 4:]
    return x


def spatial_shift2(x):
    b, w, h, c = x.size()
    x[:, :, 1:, :c // 4] = x[:, :, :h - 1, :c // 4]
    x[:, :, :h - 1, c // 4:c // 2] = x[:, :, 1:, c // 4:c // 2]
    x[:, 1:, :, c // 2:c * 3 // 4] = x[:, :w - 1, :, c // 2:c * 3 // 4]
    x[:, :w - 1, :, 3 * c // 4:] = x[:, 1:, :, 3 * c // 4:]
    return x


class SplitAttention(nn.Module):
    def __init__(self, channel=256, k=3):
        super().__init__()
        self.channel = channel
        self.k = k
        self.mlp1 = nn.Linear(channel, channel, bias=False)
        self.gelu = nn.GELU()
        self.mlp2 = nn.Linear(channel, channel * k, bias=False)
        self.softmax = nn.Softmax(1)

    def forward(self, x_all):
        b, k, h, w, c = x_all.shape
        x_all = x_all.reshape(b, k, -1, c)  # bs,k,n,c
        a = torch.sum(torch.sum(x_all, 1), 1)  # bs,c
        hat_a = self.mlp2(self.gelu(self.mlp1(a)))  # bs,kc
        hat_a = hat_a.reshape(b, self.k, c)  # bs,k,c
        bar_a = self.softmax(hat_a)  # bs,k,c
        attention = bar_a.unsqueeze(-2)  # #bs,k,1,c
        out = attention * x_all  # #bs,k,n,c
        out = torch.sum(out, 1).reshape(b, h, w, c)
        return out


class S2Attention(nn.Module):

    def __init__(self, channels=256, c2=512):
        super().__init__()
        self.mlp1 = nn.Linear(channels, channels * 3)
        self.mlp2 = nn.Linear(channels, channels)
        self.split_attention = SplitAttention(channel=channels)

    def forward(self, x):
        b, c, w, h = x.size()
        x = x.permute(0, 2, 3, 1)
        x = self.mlp1(x)
        x1 = spatial_shift1(x[:, :, :, :c])
        x2 = spatial_shift2(x[:, :, :, c:c * 2])
        x3 = x[:, :, :, c * 2:]
        x_all = torch.stack([x1, x2, x3], 1)
        a = self.split_attention(x_all)
        x = self.mlp2(a)
        x = x.permute(0, 3, 1, 2)
        return x


if __name__ == '__main__':
    input = torch.randn(1, 128, 20, 20)
    s2att = S2Attention(channels=128)
    output = s2att(input)
    print(output.shape)

注册

在ultralytics/nn/tasks.py中进行如下操作：

步骤1:

from ultralytics.nn.attention.S2Attention import S2Attention

步骤2

修改def parse_model(d, ch, verbose=True):

        if m in (
            Classify, Conv, ConvTranspose, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, Focus,
            BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, nn.ConvTranspose2d,DWConvTranspose2d, C3x, RepC3, EVCBlock,
            CAFMAttention, CloFormerAttnConv, C2f_iAFF, CSPStage, nn.Conv2d, GSConv, VoVGSCSP, C2f_CBAM, C2f_RefConv,C2f_SimAM,
            DualConv, C2f_NAM, MSFE, C2f_MDCR,C2f_MHSA, PPA, C2f_DASI, C2f_CascadedGroupAttention, C2f_DWR, C2f_DWRSeg,
            SPConv, AKConv, C2f_DySnakeConv, ScConv, C2f_ScConv, RFAConv, RFCBAMConv, RFCAConv, CAConv, CBAMConv, C2f_MultiDilatelocalAttention,
            C2f_iRMB, C2f_MSBlock, MobileViTBlock,CoordAtt, MCALayer,C2f_FocusedLinearAttention, RCSOSA,RepNCSPELAN4, SPPELAN, S2Attention
        ):
            c1, c2 = ch[f], args[0]
            if c2 != nc:  # if c2 not equal to number of classes (i.e. for Classify() output)
                c2 = make_divisible(min(c2, max_channels) * width, 8)

            args = [c1, c2, *args[1:]]
            if m in (BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, C3x, RepC3, C2f_deformable_LKA,Sea_AttentionBlock, C2f_iAFF,CSPStage, VoVGSCSP,C2f_SimAM,
                    C2f_NAM, C2f_MDCR, C2f_DASI, C2f_DWR, C2f_DWRSeg , C2f_MultiDilatelocalAttention, C2f_iRMB, C2f_MSBlock, MobileViTBlock,C2f_FocusedLinearAttention
                    ):
                args.insert(2, n)  # number of repeats
                n = 1

配置yolov8_S2Attention.yaml

ultralytics/ultralytics/cfg/models/v8/yolov8_S2Attention.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: 80 # 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]
  s: [0.33, 0.50, 1024]
  m: [0.67, 0.75, 768]
  l: [1.00, 1.00, 512]
  x: [1.00, 1.25, 512]

# 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

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

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  - [-1, 3, C2f, [256]] # 15 (P3/8-small)
  - [-1, 1, S2Attention, [256]] # 16
  
  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 12], 1, Concat, [1]] # cat head P4
  - [-1, 3, C2f, [512]] # 19 (P4/16-medium)
  - [-1, 1, S2Attention, [512]] # 20

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

  - [[16, 20, 24], 1, Detect, [nc]] # Detect(P3, P4, P5)

实验

脚本


import os
from ultralytics import YOLO
 
yaml = 'ultralytics/cfg/models/v8/yolov8_S2Attention.yaml'

 
model = YOLO(yaml) 
 
model.info()
if __name__ == "__main__":
   
    results = model.train(data='coco128.yaml',
                          name='yolov8_S2Attention',
                          epochs=10, 
                          workers=8, 
                          batch=1)

结果

YOLO大王

关注

7
点赞
踩
9

收藏

觉得还不错? 一键收藏
打赏
0
评论
【YOLOv8改进 - 注意力机制】S2Attention : 整合空间位移和分割注意力

最近，基于MLP的视觉骨干网络开始出现。与CNN和视觉Transformer相比，具有较少归纳偏差的MLP架构在图像识别中表现出竞争力。其中，采用直接空间移位操作的空间移位MLP（S2-MLP）比包括MLP-mixer和ResMLP在内的早期工作取得了更好的性能。最近，使用较小的补丁和金字塔结构，Vision Permutator（ViP）和Global Filter Network（GFNet）在性能上超过了S2-MLP。本文中，我们改进了S2-MLP视觉骨干网络。
复制链接

扫一扫