【YOLO改进】主干插入SequentialPolarizedSelfAttention模块(基于MMYOLO)

最新推荐文章于 2025-08-23 13:34:52 发布

原创

最新推荐文章于 2025-08-23 13:34:52 发布 · 782 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #深度学习 #计算机视觉

SequentialPolarizedSelfAttention模块

论文链接:https://arxiv.org/pdf/2107.00782.pdf

将SequentialPolarizedSelfAttention模块添加到MMYOLO中

将开源代码SequentialSelfAttention.py文件复制到mmyolo/models/plugins目录下
导入MMYOLO用于注册模块的包: from mmyolo.registry import MODELS
确保 class SequentialPolarizedSelfAttention中的输入维度为in_channels(因为MMYOLO会提前传入输入维度参数，所以要保持参数名的一致)
利用@MODELS.register_module()将“class SequentialPolarizedSelfAttention(nn.Module)”注册:
修改mmyolo/models/plugins/__init__.py文件
在终端运行:
```
python setup.py install
```
修改对应的配置文件，并且将plugins的参数“type”设置为“SequentialPolarizedSelfAttention”，可参考【YOLO改进】主干插入注意力机制模块CBAM（基于MMYOLO）-优快云博客

修改后的SequentialSelfAttention.py

import torch
from torch import nn
from mmyolo.registry import MODELS

@MODELS.register_module()
class SequentialPolarizedSelfAttention(nn.Module):

    def __init__(self, in_channels=512):
        super().__init__()
        self.ch_wv=nn.Conv2d(in_channels,in_channels//2,kernel_size=(1,1))
        self.ch_wq=nn.Conv2d(in_channels,1,kernel_size=(1,1))
        self.softmax_channel=nn.Softmax(1)
        self.softmax_spatial=nn.Softmax(-1)
        self.ch_wz=nn.Conv2d(in_channels//2,in_channels,kernel_size=(1,1))
        self.ln=nn.LayerNorm(in_channels)
        self.sigmoid=nn.Sigmoid()
        self.sp_wv=nn.Conv2d(in_channels,in_channels//2,kernel_size=(1,1))
        self.sp_wq=nn.Conv2d(in_channels,in_channels//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