1. Polarized Self-Attention介绍

1.1  摘要：像素级回归可能是细粒度计算机视觉任务中最常见的问题，例如估计关键点热图和分割掩模。 这些回归问题非常具有挑战性，特别是因为它们需要在低计算开销的情况下对高分辨率输入/输出的长期依赖性进行建模，以估计高度非线性的像素语义。 虽然深度卷积神经网络（DCNN）中的注意力机制在增强远程依赖性方面已变得流行，但特定于元素的注意力（例如非局部块）的学习非常复杂且对噪声敏感，并且大多数简化的注意力混合体试图达到 多种类型任务之间的最佳折衷方案。 在本文中，我们提出了偏振自注意力（PSA）模块，它结合了高质量像素级回归的两个关键设计：（1）偏振过滤：在通道和空间注意力计算中保持高内部分辨率，同时完全折叠输入张量 它们的对应尺寸。 (2) 增强：构建直接拟合典型细粒度回归输出分布的非线性，例如二维高斯分布（关键点热图）或二维二项分布（二元分割掩模）。 PSA 似乎已经耗尽了其仅通道和仅空间分支内的表示能力，因此其顺序布局和并行布局之间仅存在边际度量差异。 实验结果表明，PSA 将标准基线提高了 2−4 个点，并将 2D 姿态估计和语义分割基准的最先进技术提高了 1−2 个点。

官方论文地址：https://export.arxiv.org/pdf/2107.00782.pdf

官方代码地址：https://github.com/DeLightCMU/PSA

1.2  简单介绍:  

          Polarized Self-Attention (PSA) 模块是一种用于高质量像素级回归的自注意力机制。它主要通过两个关键设计来提高模型性能：极化滤波和增强处理。

           极化滤波：PSA在通道和空间注意力计算中保持高内部分辨率，同时沿着其对应的维度完全折叠输入张量。这有助于在不增加太多计算负担的情况下，保留重要的空间细节和通道信息。

           增强处理：PSA采用非线性组合直接拟合典型细粒度回归的输出分布，如二维高斯分布（关键点热图）或二维二项分布（二值分割掩膜）。这种设计使得PSA能够更好地适应不同任务的特定需求，并提高整体表现。

           此外，PSA模块在实验中显示出显著的性能提升，相比标准基线模型，它在2D姿态估计和语义分割基准测试中分别提高了2到4个百分点和1到2个百分点的表现。这表明PSA不仅能有效处理高分辨率信息，还能通过其独特的非线性组合进一步提高模型的预测精度。

1.3  Polarized Self-Attention模块结构图

2. 核心代码

import torch
import torch.nn as nn
 
 
class PolarizedSelfAttention(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(x)  # bs,c//2,h,w
        spatial_wq = self.sp_wq(x)  # 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 * x
        out = spatial_out + channel_out
        return out

3. YOLOv11中添加Polarized Self-Attention

3.1 在ultralytics/nn下新建Extramodule

 3.2 在Extramodule里创建PSA

在PSA.py文件里添加给出的PSA代码

添加完PSA代码后，在ultralytics/nn/Extramodule/__init__.py文件中引用

3.3 在tasks.py里引用

在ultralytics/nn/tasks.py文件里引用Extramodule

在tasks.py找到parse_model（ctrl+f可以直接搜索parse_model位置）

添加如下代码：

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

4. 新建一个yolo11PSA.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: 1 # 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, 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, PolarizedSelfAttention, []]

  - [-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, PolarizedSelfAttention, []]

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

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

  - [[17, 21, 26], 1, Detect, [nc]] # Detect(P3, P4, P5)

大家根据自己的数据集实际情况，修改nc大小。

5. 模型训练

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO(r'D:\yolo\yolov11\ultralytics-main\datasets\yolo11PSA.yaml')
    model.train(data=r'D:\yolo\yolov11\ultralytics-main\datasets\data.yaml',
                cache=False,
                imgsz=640,
                epochs=100,
                single_cls=False,  # 是否是单类别检测
                batch=4,
                close_mosaic=10,
                workers=0,
                device='0',
                optimizer='SGD',
                amp=True,
                project='runs/train',
                name='exp',
                )

模型结构打印，成功运行：

6.本文总结

到此本文的正式分享内容就结束了，在这里给大家推荐我的YOLOv11改进有效涨点专栏，本专栏目前为新开的，后期我会根据各种前沿顶会进行论文复现，也会对一些老的改进机制进行补充，如果大家觉得本文帮助到你了，订阅本专栏，关注后续更多的更新~

YOLOv11有效涨点专栏