前言

本文介绍了三重注意力机制（Triplet Attention），这是一种通过三分支结构捕获跨维度交互以计算注意力权重的轻量化方法。该方法利用旋转操作构建通道与空间维度间的依赖关系，有效编码通道间和空间信息，且计算开销极低。我们将 Triplet Attention 模块集成进 YOLOv8 的 Neck 部分，插入于特征融合后的 C2f 模块之后，增强了模型对关键特征的提取能力。实验结果表明，引入 Triplet Attention 的 YOLOv8 模型在目标检测任务中能够以可忽略的额外计算成本显著提升检测精度。

文章目录： YOLOv8改进大全：卷积层、轻量化、注意力机制、损失函数、Backbone、SPPF、Neck、检测头全方位优化汇总

专栏链接: YOLOv8改进专栏

介绍

摘要

得益于在通道或空间位置之间构建相互依赖关系的能力，注意力机制在最近被广泛研究并广泛应用于各种计算机视觉任务中。在本文中，我们研究了轻量但有效的注意力机制，并提出了三重注意力，这是一种通过使用三分支结构捕获跨维度交互来计算注意力权重的新方法。对于输入张量，三重注意力通过旋转操作及后续的残差变换构建维度间依赖关系，并以可忽略的计算开销编码通道间和空间信息。我们的方法简单且高效，可以作为附加模块轻松插入经典骨干网络中。我们在各种具有挑战性的任务中证明了我们方法的有效性，包括 ImageNet-1k 上的图像分类以及 MSCOCO 和 PASCAL VOC 数据集上的目标检测。此外，我们通过可视化检查 GradCAM 和 GradCAM++ 结果，提供了对三重注意力性能的广泛见解。我们方法的实证评估支持了在计算注意力权重时捕捉跨维度依赖关系的重要性。本文的代码可在 https://github.com/LandskapeAI/triplet-attention 公开获取。

文章链接

论文地址：论文地址

代码地址：代码地址

基本原理

核心代码

importtorchimporttorch.nnasnnclassBasicConv(nn.Module):def__init__(self,in_planes,out_planes,kernel_size,stride=1,padding=0,dilation=1,groups=1,relu=True,bn=True,bias=False,):super(BasicConv,self).__init__()self.out_channels=out_planes self.conv=nn.Conv2d(in_planes,out_planes,kernel_size=kernel_size,stride=stride,padding=padding,dilation=dilation,groups=groups,bias=bias,)self.bn=(nn.BatchNorm2d(out_planes,eps=1e-5,momentum=0.01,affine=True)ifbnelseNone)self.relu=nn.ReLU()ifreluelseNonedefforward(self,x):x=self.conv(x)ifself.bnisnotNone:x=self.bn(x)ifself.reluisnotNone:x=self.relu(x)returnxclassChannelPool(nn.Module):defforward(self,x):returntorch.cat((torch.max(x,1)[0].unsqueeze(1),torch.mean(x,1).unsqueeze(1)),dim=1)classSpatialGate(nn.Module):def__init__(self):super(SpatialGate,self).__init__()kernel_size=7self.compress=ChannelPool()self.spatial=BasicConv(2,1,kernel_size,stride=1,padding=(kernel_size-1)//2,relu=False)defforward(self,x):x_compress=self.compress(x)x_out=self.spatial(x_compress)scale=torch.sigmoid_(x_out)returnx*scaleclassTripletAttention(nn.Module):def__init__(self,gate_channels,reduction_ratio=16,pool_types=["avg","max"],no_spatial=False,):super(TripletAttention,self).__init__()self.ChannelGateH=SpatialGate()self.ChannelGateW=SpatialGate()self.no_spatial=no_spatialifnotno_spatial:self.SpatialGate=SpatialGate()defforward(self,x):x_perm1=x.permute(0,2,1,3).contiguous()x_out1=self.ChannelGateH(x_perm1)x_out11=x_out1.permute(0,2,1,3).contiguous()x_perm2=x.permute(0,3,2,1).contiguous()x_out2=self.ChannelGateW(x_perm2)x_out21=x_out2.permute(0,3,2,1).contiguous()ifnotself.no_spatial:x_out=self.SpatialGate(x)x_out=(1/3)*(x_out+x_out11+x_out21)else:x_out=(1/2)*(x_out11+x_out21)returnx_out

引入代码

在根目录下的ultralytics/nn/目录，新建一个attention目录，然后新建一个以TripletAttention为文件名的py文件， 把代码拷贝进去。

importtorchimporttorch.nnasnnclassBasicConv(nn.Module):def__init__(self,in_planes,out_planes,kernel_size,stride=1,padding=0,dilation=1,groups=1,relu=True,bn=True,bias=False):super(BasicConv,self).__init__()self.out_channels=out_planes self.conv=nn.Conv2d(in_planes,out_planes,kernel_size=kernel_size,stride=stride,padding=padding,dilation=dilation,groups=groups,bias=bias)self.bn=nn.BatchNorm2d(out_planes,eps=1e-5,momentum=0.01,affine=True)ifbnelseNoneself.relu=nn.ReLU()ifreluelseNonedefforward(self,x):x=self.conv(x)ifself.bnisnotNone:x=self.bn(x)ifself.reluisnotNone:x=self.relu(x)returnxclassZPool(nn.Module):defforward(self,x):returntorch.cat((torch.max(x,1)[0].unsqueeze(1),torch.mean(x,1).unsqueeze(1)),dim=1)classAttentionGate(nn.Module):def__init__(self):super(AttentionGate,self).__init__()kernel_size=7self.compress=ZPool()self.conv=BasicConv(2,1,kernel_size,stride=1,padding=(kernel_size-1)//2,relu=False)defforward(self,x):x_compress=self.compress(x)x_out=self.conv(x_compress)scale=torch.sigmoid_(x_out)returnx*scaleclassTripletAttention(nn.Module):def__init__(self,no_spatial=False):super(TripletAttention,self).__init__()self.cw=AttentionGate()self.hc=AttentionGate()self.no_spatial=no_spatialifnotno_spatial:self.hw=AttentionGate()defforward(self,x):x_perm1=x.permute(0,2,1,3).contiguous()x_out1=self.cw(x_perm1)x_out11=x_out1.permute(0,2,1,3).contiguous()x_perm2=x.permute(0,3,2,1).contiguous()x_out2=self.hc(x_perm2)x_out21=x_out2.permute(0,3,2,1).contiguous()ifnotself.no_spatial:x_out=self.hw(x)x_out=1/3*(x_out+x_out11+x_out21)else:x_out=1/2*(x_out11+x_out21)returnx_out

注册

在ultralytics/nn/tasks.py中进行如下操作：

步骤1:

fromultralytics.nn.attention.TripletAttentionimportTripletAttention

步骤2

修改def parse_model(d, ch, verbose=True):

elifmin{TripletAttention}:c2=ch[f]args=[c2,*args]

配置yolov8_TripletAttention.yaml

ultralytics/cfg/models/v8/yolov8_TripletAttention.yaml

# Ultralytics YOLO 🚀, AGPL-3.0 license# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect# Parametersnc:2# number of classesscales:# model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'# [depth, width, max_channels]n:[0.33,0.25,1024]# YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPss:[0.33,0.50,1024]# YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPsm:[0.67,0.75,768]# YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPsl:[1.00,1.00,512]# YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPsx:[1.00,1.25,512]# YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs# YOLOv8.0n backbonebackbone:# [from, repeats, module, args]-[-1,1,Conv,[64,3,2]]# 0-P1/2-[-1,1,Conv,[128,3,2]]# 1-P2/4-[-1,3,C2f,[128,True]]-[-1,1,Conv,[256,3,2]]# 3-P3/8-[-1,6,C2f,[256,True]]-[-1,1,Conv,[512,3,2]]# 5-P4/16-[-1,6,C2f,[512,True]]-[-1,1,Conv,[1024,3,2]]# 7-P5/32-[-1,3,C2f,[1024,True]]-[-1,1,SPPF,[1024,5]]# 9# YOLOv8.0n headhead:-[-1,1,nn.Upsample,[None,2,'nearest']]-[[-1,6],1,Concat,[1]]# cat backbone P4-[-1,3,C2f,[512]]# 12-[-1,1,nn.Upsample,[None,2,'nearest']]-[[-1,4],1,Concat,[1]]# cat backbone P3-[-1,3,C2f,[256]]# 15 (P3/8-small)-[-1,1,TripletAttention,[]]# 16-[-1,1,Conv,[256,3,2]]-[[-1,12],1,Concat,[1]]# cat head P4-[-1,3,C2f,[512]]# 19 (P4/16-medium)-[-1,1,TripletAttention,[]]# 20-[-1,1,Conv,[512,3,2]]-[[-1,9],1,Concat,[1]]# cat head P5-[-1,3,C2f,[1024]]# 23 (P5/32-large)-[-1,1,TripletAttention,[]]# 24-[[16,20,24],1,Detect,[nc]]# Detect(P3, P4, P5)

实验

脚本

importosfromultralyticsimportYOLO# Define the configuration options directlyyaml='ultralytics/cfg/models/v8/yolov8_TripletAttention.yaml'# Initialize the YOLO model with the specified YAML filemodel=YOLO(yaml)# Print model informationmodel.info()if__name__=="__main__":# Train the model with the specified parametersresults=model.train(data='ultralytics/datasets/original-license-plates.yaml',name='yolov8_TripletAttention',epochs=10,workers=8,batch=1)

结果