YOLOv5改进 | 添加注意力篇 | 利用YOLOv10提出的PSA注意力机制助力YOLOv5有效涨点（附代码 + 详细修改教程）

一、本文介绍

本文给大家带来的改进机制是YOLOv10提出的PSA注意力机制，自注意力在各种视觉任务中得到了广泛应用，因为它具有显著的全局建模能力。然而，自注意力机制表现出较高的计算复杂度和内存占用。为了解决这个问题，鉴于注意力头冗余的普遍存在，我们提出了一种高效的部分自注意力（PSA）模块设计，其能够在不显著增加计算成本的情况下提升YOLO模型的性能！本文附其网络结构图辅助大家理解该结构，同时本文包含YOLOv8添加该注意力机制的方法！

欢迎大家订阅我的专栏一起学习YOLO！

专栏目录：YOLOv5改进有效涨点目录 | 包含卷积、主干、检测头、注意力机制、Neck上百种创新机制

一、本文介绍

二、PSA介绍

三、PSA核心代码

四、手把手教你添加PSA注意力机制

4.1 PSA添加步骤

4.1.1 修改一

4.1.2 修改二

4.1.3 修改三

4.1.4 修改四

4.2 PSA的yaml文件

4.3 训练截图

五、本文总结

二、PSA介绍

论文地址：官方论文地址点击此处即可跳转

代码地址：官方代码地址点击此处即可跳转

自注意力在各种视觉任务中得到了广泛应用，因为它具有显著的全局建模能力。然而，自注意力机制表现出较高的计算复杂度和内存占用。为了解决这个问题，鉴于注意力头冗余的普遍存在，我们提出了一种高效的部分自注意力（PSA）模块设计，如图3(c)所示。

具体来说，我们通过1×1卷积将特征均匀地划分为两部分。然后，我们仅将其中一部分输入到由多头自注意力模块（MHSA）和前馈网络（FFN）组成的NPSA块中。两部分特征随后被连接并通过1×1卷积融合。此外，我们遵循将MHSA中查询和键的维度分配为值的一半，并将LayerNorm替换为BatchNorm以加快推理速度。

PSA仅在分辨率最低的Stage 4之后放置，以避免自注意力二次复杂性带来的过多开销。通过这种方式，可以在低计算成本下将全局表示学习能力引入YOLO模型，从而增强模型能力并提高性能。

通过这些精度驱动的设计，我们能够在不显著增加计算成本的情况下提升YOLO模型的性能。

三、PSA核心代码

核心代码的使用方式看章节四！

import torch
import torch.nn as nn__all__ = ['PSA']def autopad(k, p=None, d=1):  # kernel, padding, dilation"""Pad to 'same' shape outputs."""if d > 1:k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-sizeif p is None:p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-padreturn pclass Conv(nn.Module):"""Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)."""default_act = nn.SiLU()  # default activationdef __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):"""Initialize Conv layer with given arguments including activation."""super().__init__()self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)self.bn = nn.BatchNorm2d(c2)self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()def forward(self, x):"""Apply convolution, batch normalization and activation to input tensor."""return self.act(self.bn(self.conv(x)))def forward_fuse(self, x):"""Perform transposed convolution of 2D data."""return self.act(self.conv(x))class Attention(nn.Module):def __init__(self, dim, num_heads=8,attn_ratio=0.5):super().__init__()self.num_heads = num_headsself.head_dim = dim // num_headsself.key_dim = int(self.head_dim * attn_ratio)self.scale = self.key_dim ** -0.5nh_kd = nh_kd = self.key_dim * num_headsh = dim + nh_kd * 2self.qkv = Conv(dim, h, 1, act=False)self.proj = Conv(dim, dim, 1, act=False)self.pe = Conv(dim, dim, 3, 1, g=dim, act=False)def forward(self, x):B, _, H, W = x.shapeN = H * Wqkv = self.qkv(x)q, k, v = qkv.view(B, self.num_heads, -1, N).split([self.key_dim, self.key_dim, self.head_dim], dim=2)attn = ((q.transpose(-2, -1) @ k) * self.scale)attn = attn.softmax(dim=-1)x = (v @ attn.transpose(-2, -1)).view(B, -1, H, W) + self.pe(v.reshape(B, -1, H, W))x = self.proj(x)return xclass PSA(nn.Module):def __init__(self, c1, c2, e=0.5):super().__init__()assert (c1 == c2)self.c = int(c1 * e)self.cv1 = Conv(c1, 2 * self.c, 1, 1)self.cv2 = Conv(2 * self.c, c1, 1)self.attn = Attention(self.c, attn_ratio=0.5, num_heads=self.c // 64)self.ffn = nn.Sequential(Conv(self.c, self.c * 2, 1),Conv(self.c * 2, self.c, 1, act=False))def forward(self, x):a, b = self.cv1(x).split((self.c, self.c), dim=1)b = b + self.attn(b)b = b + self.ffn(b)return self.cv2(torch.cat((a, b), 1))

四、手把手教你添加PSA注意力机制

4.1 PSA添加步骤

4.1.1 修改一

首先我们找到如下的目录'yolov5-master/models'，然后在这个目录下在创建一个新的目录然后这个就是存储改进的仓库，大家可以在这里新建所有的改进的py文件，对应改进的文件名字可以根据你自己的习惯起(不影响任何但是下面导入的时候记住改成你对应的即可)，然后将PSA的核心代码复制进去。

4.1.2 修改二

然后在新建的目录里面我们在新建一个__init__.py文件(此文件大家只需要建立一个即可)，然后我们在里面添加导入我们模块的代码。注意标记一个'.'其作用是标记当前目录。

4.1.3 修改三

然后我们找到如下文件''models/yolo.py''在开头的地方导入我们的模块按照如下修改->

(如果你看了我多个改进机制此处只需要添加一个即可，无需重复添加)

注意的添加位置要放在common的导入上面！！！！！

4.1.4 修改四

然后我们找到parse_model方法，按照如下修改->

到此就修改完成了，复制下面的ymal文件即可运行。

4.2 PSA的yaml文件

保持和YOLOv10相同的使用方法！

# YOLOv5 v6.0 backbone
backbone:# [from, number, module, args][[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2[-1, 1, Conv, [128, 3, 2]], # 1-P2/4[-1, 3, C3, [128]],[-1, 1, Conv, [256, 3, 2]], # 3-P3/8[-1, 6, C3, [256]],[-1, 1, Conv, [512, 3, 2]], # 5-P4/16[-1, 9, C3, [512]],[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32[-1, 3, C3, [1024]],[-1, 1, SPPF, [1024, 5]], # 9[-1, 1, PSA, [1024]],  # 10]# YOLOv5 v6.0 head
head: [[-1, 1, Conv, [512, 1, 1]],[-1, 1, nn.Upsample, [None, 2, "nearest"]],[[-1, 6], 1, Concat, [1]], # cat backbone P4[-1, 3, C3, [512, False]], # 14[-1, 1, Conv, [256, 1, 1]],[-1, 1, nn.Upsample, [None, 2, "nearest"]],[[-1, 4], 1, Concat, [1]], # cat backbone P3[-1, 3, C3, [256, False]], # 18 (P3/8-small)[-1, 1, Conv, [256, 3, 2]],[[-1, 15], 1, Concat, [1]], # cat head P4[-1, 3, C3, [512, False]], # 21 (P4/16-medium)[-1, 1, Conv, [512, 3, 2]],[[-1, 11], 1, Concat, [1]], # cat head P5[-1, 3, C3, [1024, False]], # 24 (P5/32-large)[[18, 21, 24], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)]