YOLOv9改进策略【注意力机制篇】| PSA极化自我关注： 实现高质量像素回归

一、本文介绍

本文记录的是基于PSA注意力模块的YOLOv9目标检测方法研究。PSA模块通过极化滤波和增强设计，提高了内部分辨率，并增强非线性拟合，从而能够提升像素级回归任务的性能。本文将其应用到YOLOv9的检测任务中，使模型能够更好地捕捉图像中的细节信息，以实现目标检测任务中准确识别和定位。

文章目录

• 一、本文介绍
• 二、PSA注意力原理
• • 2.1、设计原理
  • 2.2、优势
• 三、PSA的实现代码
• 四、添加步骤
• • 4.1 修改common.py
  • • 4.1.1 基础模块1
    • 4.1.2 创新模块2⭐
  • 4.2 修改yolo.py
• 五、yaml模型文件
• • 5.1 模型改进版本一
  • 5.2 模型改进版本二⭐
• 六、成功运行结果

二、PSA注意力原理

极化自我关注： 实现高质量像素回归

PSA（Polarized Self - Attention）模块是一种用于高质量像素级回归的注意力模块，其设计原理和优势主要包括以下几个方面：

2.1、设计原理

• 极化滤波（Polarized filtering）：PSA在通道和空间注意力计算中保持高内部分辨率，同时在输入张量的对应维度上完全折叠。具体来说，通过通道 - 只关注分支（Channel - only branch）和空间 - 只关注分支（Spatial - only branch）来实现。
  • 通道 - 只关注分支： A c h ( X ) = F S G [ W z ∣ θ 1 ( ( σ 1 ( W v ( X ) ) × F S M ( σ 2 ( W q ( X ) ) ) ) ) ] A^{ch}(X)=F_{SG}[W_{z|\theta_{1}}((\sigma_{1}(W_{v}(X)) \times F_{SM}(\sigma_{2}(W_{q}(X)))))] Ach(X)=FSG​[Wz∣θ1​​((σ1​(Wv​(X))×FSM​(σ2​(Wq​(X)))))]，其中 W q W_{q} Wq​、 W v W_{v} Wv​和 w z w_{z} wz​是 1 × 1 1\times1 1×1卷积层， σ 1 \sigma_{1} σ1​和 σ 2 \sigma_{2} σ2​是张量重塑操作， F S M F_{SM} FSM​是SoftMax算子，内部通道数在 W v W_{v} Wv​、 W q W_{q} Wq​和 w z w_{z} wz​之间为 C / 2 C/2 C/2。输出为 Z c h = A c h ( X ) ⊙ c h X Z^{ch}=A^{ch}(X)\odot^{ch}X Zch=Ach(X)⊙chX，其中 ⊙ c h \odot^{ch} ⊙ch是通道乘法运算符。
  • 空间 - 只关注分支： A s p ( X ) = F S G [ σ 3 ( F S M ( σ 1 ( F G P ( W q ( X ) ) ) ) × σ 2 ( W v ( X ) ) ) ] A^{sp}(X)=F_{SG}[\sigma_{3}(F_{SM}(\sigma_{1}(F_{GP}(W_{q}(X)))) \times \sigma_{2}(W_{v}(X)))] Asp(X)=FSG​[σ3​(FSM​(σ1​(FGP​(Wq​(X))))×σ2​(Wv​(X)))]，其中 w q w_{q} wq​和 w v w_{v} wv​是标准 1 × 1 1\times1 1×1卷积层， θ 2 \theta_{2} θ2​是中间参数， σ 1 \sigma_{1} σ1​、 σ 2 \sigma_{2} σ2​和 σ 3 \sigma_{3} σ3​是张量重塑操作， F S M F_{SM} FSM​是SoftMax算子， F G P F_{GP} FGP​是全局池化算子。输出为 Z s p = A s p ( X ) ⊙ s p X Z^{sp}=A^{sp}(X)\odot^{sp}X Zsp=Asp(X)⊙spX，其中 ⊙ s p \odot^{sp} ⊙sp是空间乘法运算符。
• 增强（Enhancement）：在通道 - 只关注和空间 - 只关注分支中融合了非线性，以直接拟合典型细粒度回归的输出分布，例如2D高斯分布（关键点热图）或2D二项分布（二值分割掩码）。具体来说，使用了Softmax-Sigmoid组合。
• 组合（Composition）：两个分支的输出可以在并行布局 P S A p ( X ) = Z c h + Z s p = A c h ( X ) ⊙ c h X + A s p ( X ) ⊙ s p X PSA_{p}(X)=Z^{ch} + Z^{sp}=A^{ch}(X)\odot^{ch}X + A^{sp}(X)\odot^{sp}X PSAp​(X)=Zch+Zsp=Ach(X)⊙chX+Asp(X)⊙spX或顺序布局 P S A s ( X ) = Z s p ( Z c h ) = A s p ( A c h ( X ) ⊙ c h X ) ⊙ s p A c h ( X ) ⊙ c h X PSA_{s}(X)=Z^{sp}(Z^{ch})=A^{sp}(A^{ch}(X)\odot^{ch}X)\odot^{sp}A^{ch}(X)\odot^{ch}X PSAs​(X)=Zsp(Zch)=Asp(Ach(X)⊙chX)⊙spAch(X)⊙chX下进行组合。

2.2、优势

• 保持高内部分辨率：与现有注意力模块在其最佳配置下相比，PSA在通道（ C / 2 C/2 C/2）和空间（ [ W , H ] [W, H] [W,H]）维度上保留了最高的注意力分辨率，有助于保留高分辨率信息。
• 高效的参数利用：
  • 通道 - 只关注分支：在通道 - 只关注注意力中，Softmax重加权与挤压 - 激励相结合，利用Softmax作为大小为 C / 2 × W × H C/2\times W\times H C/2×W×H的瓶颈张量的非线性激活，遵循了受益于GC和SE块的挤压 - 激励模式，在可比的GC块计算复杂度下进行更高分辨率的挤压 - 激励。
  • 空间 - 只关注分支：不仅保持了完整的 [ W , H ] [W, H] [W,H]空间分辨率，而且在内部为非线性Softmax重加权保留了 2 × C × C / 2 2\times C\times C/2 2×C×C/2个可学习参数，比现有模块更强大。
• 拟合输出分布：PSA的通道 - 只关注和空间 - 只关注分支都使用Softmax - Sigmoid组合，能够近似多模式高斯图（关键点热图）和分段二项图（分割掩码），充分利用PSA注意力分支中保留的高分辨率信息。
• 表示能力强：PSA的并行和顺序布局之间的指标差异较小，这表明其通道 - 只关注和空间 - 只关注注意力块的设计可能已经充分利用了沿通道和空间维度的表示能力。

论文：https://arxiv.org/pdf/2107.00782
源码：https://github.com/DeLightCMU/PSA

三、PSA的实现代码

PSA模块的实现代码如下：

class PSA_Channel(nn.Module):
    def __init__(self, c1) -> None:
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = nn.Conv2d(c1, c_, 1)
        self.cv2 = nn.Conv2d(c1, 1, 1)
        self.cv3 = nn.Conv2d(c_, c1, 1)
        self.reshape1 = nn.Flatten(start_dim=-2, end_dim=-1)
        self.reshape2 = nn.Flatten()
        self.sigmoid = nn.Sigmoid()
        self.softmax = nn.Softmax(1)
        self.layernorm = nn.LayerNorm([c1, 1, 1])

    def forward(self, x): # shape(batch, channel, height, width)
        x1 = self.reshape1(self.cv1(x)) # shape(batch, channel/2, height*width)
        x2 = self.softmax(self.reshape2(self.cv2(x))) # shape(batch, height*width)
        y = torch.matmul(x1, x2.unsqueeze(-1)).unsqueeze(-1) # 高维度下的矩阵乘法（最后两个维度相乘）
        return self.sigmoid(self.layernorm(self.cv3(y))) * x

class PSA_Spatial(nn.Module):
    def __init__(self, c1) -> None:
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = nn.Conv2d(c1, c_, 1)
        self.cv2 = nn.Conv2d(c1, c_, 1)
        self.reshape1 = nn.Flatten(start_dim=-2, end_dim=-1)
        self.globalPooling = nn.AdaptiveAvgPool2d(1)
        self.softmax = nn.Softmax(1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x): # shape(batch, channel, height, width)
        x1 = self.reshape1(self.cv1(x)) # shape(batch, channel/2, height*width)
        x2 = self.softmax(self.globalPooling(self.cv2(x)).squeeze(-1)) # shape(batch, channel/2, 1)
        y = torch.bmm(x2.permute(0,2,1), x1) # shape(batch, 1, height*width)
        return self.sigmoid(y.view(x.shape[0], 1, x.shape[2], x.shape[3])) * x

class PSA(nn.Module):
    def __init__(self, in_channel, parallel=True) -> None:
        super().__init__()
        self.parallel = parallel
        self.channel = PSA_Channel(in_channel)
        self.spatial = PSA_Spatial(in_channel)

    def forward(self, x):
        if(self.parallel):
            return self.channel(x) + self.spatial(x)
        return self.spatial(self.channel(x))

四、添加步骤

4.1 修改common.py

此处需要修改的文件是models/common.py

common.py中定义了网络结构的通用模块，我们想要加入新的模块就只需要将模块代码放到这个文件内即可。

4.1.1 基础模块1

模块改进方法1️⃣：直接加入ParallelPolarizedSelfAttention模块。
ParallelPolarizedSelfAttention模块添加后如下：

注意❗：在4.2小节中的yolo.py文件中需要声明的模块名称为：ParallelPolarizedSelfAttention。

4.1.2 创新模块2⭐

模块改进方法2️⃣：基于ParallelPolarizedSelfAttention模块的RepNCSPELAN4。

相较方法一中的直接插入注意力模块，利用注意力模块对卷积等其他模块进行改进，其新颖程度会更高一些，训练精度可能会表现的更高。

第二种改进方法是对YOLOv9中的RepNCSPELAN4模块进行改进。将PSA注意力模块替换RepNCSPELAN4中的卷积模块，目的是通过PSA的像素级回归能力，提高RepNCSPELAN4模块的特征提取性能，使模型能够更好地捕捉图像中的细节信息，以实现目标检测任务中准确识别和定位。

改进代码如下：

class PSARepNCSPELAN4(nn.Module):
    # csp-elan
    def __init__(self, c1, c2, c3, c4, c5=1):  # ch_in, ch_out, number, shortcut, groups, expansion
        super().__init__()
        self.c = c3//2
        self.cv1 = Conv(c1, c3, 1, 1)
        self.cv2 = nn.Sequential(RepNCSP(c3//2, c4, c5), ParallelPolarizedSelfAttention(c4))
        self.cv3 = nn.Sequential(RepNCSP(c4, c4, c5), ParallelPolarizedSelfAttention(c4))
        self.cv4 = Conv(c3+(2*c4), c2, 1, 1)

    def forward(self, x):
        y = list(self.cv1(x).chunk(2, 1))
        y.extend((m(y[-1])) for m in [self.cv2, self.cv3])
        return self.cv4(torch.cat(y, 1))

    def forward_split(self, x):
        y = list(self.cv1(x).split((self.c, self.c), 1))
        y.extend(m(y[-1]) for m in [self.cv2, self.cv3])
        return self.cv4(torch.cat(y, 1))

注意❗：在4.2小节中的yolo.py文件中需要声明的模块名称为：PSARepNCSPELAN4。

4.2 修改yolo.py

此处需要修改的文件是models/yolo.py

yolo.py用于函数调用，我们只需要将common.py中定义的新的模块名添加到parse_model函数下即可。

ParallelPolarizedSelfAttention模块以及PSARepNCSPELAN4模块添加后如下：

五、yaml模型文件

5.1 模型改进版本一

在代码配置完成后，配置模型的YAML文件。

此处以models/detect/yolov9-c.yaml为例，在同目录下创建一个用于自己数据集训练的模型文件yolov9-c-psa.yaml。

将yolov9-c.yaml中的内容复制到yolov9-c-psa.yaml文件下，修改nc数量等于自己数据中目标的数量。
在骨干网络的最后一层添加psa模块，即下方代码中的第45行，只需要填入一个参数，通道数，和前一层通道数一致。