YOLOv8改进 | 注意力机制 | 迈向高质量像素级回归的极化自注意力【全网独家】

秋招面试专栏推荐 ：深度学习算法工程师面试问题总结【百面算法工程师】——点击即可跳转

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

专栏目录：《YOLOv8改进有效涨点》专栏介绍 & 专栏目录 | 目前已有40+篇内容，内含各种Head检测头、损失函数Loss、Backbone、Neck、NMS等创新点改进——点击即可跳转

虽然深度卷积神经网络中的注意力机制已经变得流行，用于增强长距离依赖，但元素特定的注意力，如Nonlocal块，学习起来非常复杂且对噪声敏感，而且大多数简化的注意力混合方法试图在多种类型的任务之间达到最佳的折衷。为此研究人员提出了极化自注意力（PSA）块，它结合了两个关键设计，以实现高质量的像素级回归。文章在介绍主要的原理后，将手把手教学如何进行模块的代码添加和修改，并将修改后的完整代码放在文章的最后，方便大家一键运行，小白也可轻松上手实践。以帮助您更好地学习深度学习目标检测YOLO系列的挑战。

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

1.原理

2. 将极化自注意力添加到YOLOv8中

2.1 PSA代码实现

2.2 更改init.py文件

2.3 添加yaml文件

2.4 在task.py中进行注册

2.5 执行程序

3. 完整代码分享

4. GFLOPs

5. 进阶

6. 总结

1.原理

论文地址：Polarized Self-Attention: Towards High-quality Pixel-wise Regression——点击即可跳转

官方代码：官方代码仓库——点击即可跳转

极化自注意力（Polarized Self-Attention，PSA）是一种针对像素级回归任务设计的注意力机制，旨在提高模型在关键点估计和语义分割等任务上的性能。其主要原理可以概括为以下几点：

1. 极化过滤：

PSA 将注意力机制应用于输入特征图，将其分为两个分支：通道分支和空间分支。
在通道分支中，PSA 沿着通道维度保持高分辨率，同时将特征图在空间维度上完全折叠，从而减少了计算量和内存占用。
在空间分支中，PSA 沿着空间维度保持高分辨率，同时将特征图在通道维度上完全折叠。
这种“极化”的设计方式有效地保留了高分辨率信息，避免了传统自注意力机制中由于池化或下采样造成的分辨率损失。

2. 非线性增强：

PSA 在通道分支和空间分支中分别使用了 Softmax 和 Sigmoid 函数的组合来增强非线性。
Softmax 函数用于将通道分支中的特征图转化为概率分布，使其更适合表示高斯分布（例如关键点热图）。
Sigmoid 函数用于将空间分支中的特征图转化为二值分布，使其更适合表示二项分布（例如分割掩码）。
这种非线性组合的设计能够更好地拟合像素级回归任务的输出分布，从而提高模型的预测精度。

3. 布局灵活：

PSA 支持两种布局方式：并行和串行。
在并行布局中，通道分支和空间分支的结果直接相加。
在串行布局中，通道分支的结果首先通过空间分支，然后两者相加。
实验表明，两种布局方式在性能上没有显著差异，这表明 PSA 的设计已经充分挖掘了通道和空间维度上的信息。

4. 优势：

PSA 在不显著增加计算量和内存占用的情况下，能够显著提高像素级回归任务的性能。
PSA 的设计能够更好地拟合像素级回归任务的输出分布，从而提高模型的预测精度。
PSA 的布局方式灵活，可以根据具体任务进行调整。

总结： PSA 通过“极化”的设计方式，有效地保留了高分辨率信息，并通过非线性增强来拟合像素级回归任务的输出分布，从而显著提高了模型的性能。

2. 将极化自注意力添加到YOLOv8中

2.1 PSA代码实现

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

class PSA(nn.Module):

    def __init__(self, channel=512):
        super().__init__()
        self.ch_wv = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))
        self.ch_wq = nn.Conv2d(channel, 1, kernel_size=(1, 1))
        self.softmax_channel = nn.Softmax(1)
        self.softmax_spatial = nn.Softmax(-1)
        self.ch_wz = nn.Conv2d(channel // 2, channel, kernel_size=(1, 1))
        self.ln = nn.LayerNorm(channel)
        self.sigmoid = nn.Sigmoid()
        self.sp_wv = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))
        self.sp_wq = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))
        self.agp = nn.AdaptiveAvgPool2d((1, 1))

    def forward(self, x):
        b, c, h, w = x.size()

        # Channel-only Self-Attention
        channel_wv = self.ch_wv(x)  # bs,c//2,h,w
        channel_wq = self.ch_wq(x)  # bs,1,h,w
        channel_wv = channel_wv.reshape(b, c // 2, -1)  # bs,c//2,h*w
        channel_wq = channel_wq.reshape(b, -1, 1)  # bs,h*w,1
        channel_wq = self.softmax_channel(channel_wq)
        channel_wz = torch.matmul(channel_wv, channel_wq).unsqueeze(-1)  # bs,c//2,1,1
        channel_weight = self.sigmoid(self.ln(self.ch_wz(channel_wz).reshape(b, c, 1).permute(0, 2, 1))).permute(0, 2,
                                                                                                                 1).reshape(
            b, c, 1, 1)  # bs,c,1,1
        channel_out = channel_weight * x

        # Spatial-only Self-Attention
        spatial_wv = self.sp_wv(channel_out)  # bs,c//2,h,w
        spatial_wq = self.sp_wq(channel_out)  # bs,c//2,h,w
        spatial_wq = self.agp(spatial_wq)  # bs,c//2,1,1
        spatial_wv = spatial_wv.reshape(b, c // 2, -1)  # bs,c//2,h*w
        spatial_wq = spatial_wq.permute(0, 2, 3, 1).reshape(b, 1, c // 2)  # bs,1,c//2
        spatial_wq = self.softmax_spatial(spatial_wq)
        spatial_wz = torch.matmul(spatial_wq, spatial_wv)  # bs,1,h*w
        spatial_weight = self.sigmoid(spatial_wz.reshape(b, 1, h, w))  # bs,1,h,w
        spatial_out = spatial_weight * channel_out
        return spatial_out

PSA (Polarized Self-Attention) 是一种针对像素级回归任务设计的注意力机制，其主要流程可以分为以下几个步骤：

1. 特征提取：

首先，使用卷积神经网络（CNN）对输入图片进行特征提取，得到特征图。特征图包含了图片中每个像素点的特征信息。

2. 极化自注意力：

将特征图输入到 PSA 模块进行特征增强。
PSA 模块包含两个分支：通道分支和空间分支。
- 通道分支：沿着通道维度保持高分辨率，同时将特征图在空间维度上完全折叠。
- 空间分支：沿着空间维度保持高分辨率，同时将特征图在通道维度上完全折叠。
在通道分支和空间分支中，分别使用 Softmax 和 Sigmoid 函数的组合来增强非线性，从而更好地拟合像素级回归任务的输出分布。

3. 输出：

将 PSA 模块处理后的特征图输入到解码器中，进行进一步的预测。
解码器会根据具体的任务类型（例如关键点估计或语义分割）生成相应的输出，例如关键点热图或分割掩码。

4. 优化：

使用损失函数来评估模型预测结果与真实标签之间的差距。
通过反向传播算法来更新模型参数，从而优化模型的性能。

总结： PSA 的主要流程可以概括为：特征提取 -> 极化自注意力 -> 输出 -> 优化。通过引入 PSA 模块，可以有效地提高模型在像素级回归任务上的性能。

2.2 更改init.py文件

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

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

2.3 添加yaml文件

关键步骤三：在/ultralytics/ultralytics/cfg/models/v8下面新建文件yolov8_PSA.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 ]  # YOLOv8n summary: 225 layers,  3157200 parameters,  3157184 gradients,   8.9 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

# 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, Conv, [ 256, 3, 2 ] ]
  - [ [ -1, 12 ], 1, Concat, [ 1 ] ]  # cat head P4
  - [ -1, 3, C2f, [ 512 ] ]  # 18 (P4/16-medium)
  - [ -1, 1, PSA, [ 512 ] ]

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

  - [ [ 15, 19, 23 ], 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函数中进行注册，

elif m is PSA :
            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, *args[1:]]

2.5 执行程序

关键步骤五：在ultralytics文件中新建train.py，将model的参数路径设置为yolov8_PSA.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_PSA.yaml')  # build from YAML and transfer weights
 
# Train the model
model.train(device = [3], batch=16)

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

3. 完整代码分享

https://pan.baidu.com/s/1bNQG8yuv5vVIbvJAMTWZ7g?pwd=n8qy

提取码: n8qy

4. GFLOPs

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

未改进的YOLOv8nGFLOPs

改进后的GFLOPs

5. 进阶

可以结合损失函数或者注意力机制进行多重改进

6. 总结

PSA 是一种针对像素级回归任务设计的注意力机制，其核心思想是通过“极化”的方式，在保持高分辨率信息的同时，有效地增强特征表示，从而提高模型的预测精度。具体而言，PSA 将注意力机制应用于输入特征图，将其分为通道分支和空间分支。在通道分支中，PSA 沿着通道维度保持高分辨率，同时将特征图在空间维度上完全折叠；在空间分支中，PSA 沿着空间维度保持高分辨率，同时将特征图在通道维度上完全折叠。这种“极化”的设计方式有效地保留了高分辨率信息，避免了传统自注意力机制中由于池化或下采样造成的分辨率损失。此外，PSA 在通道分支和空间分支中分别使用了 Softmax 和 Sigmoid 函数的组合来增强非线性，从而更好地拟合像素级回归任务的输出分布。这种非线性组合的设计能够更好地拟合像素级回归任务的输出分布，例如高斯分布（例如关键点热图）或二项分布（例如分割掩码），从而提高模型的预测精度。最后，PSA 支持两种布局方式：并行和串行，两种布局方式在性能上没有显著差异，这表明 PSA 的设计已经充分挖掘了通道和空间维度上的信息。