yolov8添加自注意力机制（Self-Attention）

### 如何在YOLOv8中添加自注意力机制 #### 添加 Polarized Self-Attention (PSA) 为了提升YOLOv8的目标检测性能，在`PolarizedSelfAttention.py`文件里加入Polarized Self-Attention(PSA)代码[^2]。该模块融合了空间和通道两个维度的信息处理能力，相较于其他单一维度的注意力机制有着更好的表现。 具体来说，PSA通过设计Channel-only分支和Spatial-only分支来分别捕捉特征图上的通道依赖性和位置关系，最终将两者结合起来形成更丰富的表征形式[^3]。 以下是Python代码片段用于定义Polarized Self-Attention层： ```python import torch.nn as nn import torch class ChannelOnlyBranch(nn.Module): """Channel-only branch of PSA.""" def __init__(self, in_channels, reduction_ratio=16): super(ChannelOnlyBranch, self).__init__() hidden_channels = int(in_channels / reduction_ratio) self.fc = nn.Sequential( nn.Linear(in_channels, hidden_channels), nn.ReLU(inplace=True), nn.Linear(hidden_channels, in_channels), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() y = self.fc(x.mean(dim=(2, 3)).view(b, c)) return x * y.view(b, c, 1, 1) class SpatialOnlyBranch(nn.Module): """Spatial-only branch of PSA.""" def __init__(self): super(SpatialOnlyBranch, self).__init__() self.conv = nn.Conv2d(2, 1, kernel_size=7, padding=3, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) scale = self.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1))) return x * scale class PolarizedSelfAttention(nn.Module): def __init__(self, channel): super(PolarizedSelfAttention, self).__init__() self.channel_branch = ChannelOnlyBranch(channel) self.spatial_branch = SpatialOnlyBranch() def forward(self, x): chan_att_x = self.channel_branch(x) spa_att_x = self.spatial_branch(chan_att_x) return spa_att_x ``` 完成上述类定义之后，需修改YOLOv8网络结构配置文件（通常是`.yaml`），以便在网络特定层次上插入此新组件。通常建议将其放置于骨干网结束处或是颈部部分之前的位置，从而让整个模型受益于增强后的特征表达力。 最后一步是在训练过程中加载并应用这些改动过的架构设置，确保所有必要的库已正确安装以及环境变量已被适当调整以支持新的特性实现。