012、ShuffleAttention与TripletAttention：轻量级注意力模块的对比与代码集成

012、ShuffleAttention与TripletAttention轻量级注意力模块的对比与代码集成上周调一个YOLOv8n的检测模型在嵌入式设备上跑发现参数量倒是压下来了但mAP掉了将近3个点。翻来覆去改backbone、调anchor效果都不理想。后来同事提了一句“你试试在Neck后面加个轻量级注意力别用SE那种带全连接的参数量上去了回不来。”这才开始认真对比ShuffleAttention和TripletAttention这两个轻量化注意力模块。先说结论ShuffleAttention适合做通道间的特征重标定TripletAttention更适合做空间位置上的特征增强。两个模块加进去参数量增加都在0.1M以内但提升的点位不一样。ShuffleAttention把通道分组打散再算注意力这个模块的思路挺有意思。传统SE注意力是全连接层压缩通道ShuffleAttention借鉴了ShuffleNet的分组思想——先把输入特征图按通道分组每组内部独立计算注意力最后通过通道混洗channel shuffle让组间信息流通。看代码实现会更直观classShuffleAttention(nn.Module):def__init__(self,channels,groups8):super().__init__()self.groupsgroups# 这里踩过坑groups必须能被channels整除否则后面reshape会报错assertchannels%groups0,fchannels{channels}must be divisible by groups{groups}self.avg_poolnn.AdaptiveAvgPool2d(1)self.max_poolnn.AdaptiveMaxPool2d(1)# 别这样写用两个独立的卷积层参数量翻倍# 正确做法共享一个卷积层分别处理avg和max的结果self.shared_convnn.Conv2d(2,1,kernel_size1,biasFalse)self.sigmoidnn.Sigmoid()defforward(self,x):b,c,h,wx.shape# 分组把通道维度拆成 groups 份xx.view(b*self.groups,-1,h,w)# 每组独立处理# 这里注意分组后每组通道数 c // groups# 通道注意力分支avg_outself.avg_pool(x)max_outself.max_pool(x)# 拼接后过卷积别写成两个独立卷积channel_atttorch.cat([avg_out,max_out],dim1)channel_attself.shared_conv(channel_att)channel_attself.sigmoid(channel_att)# 空间注意力分支spatial_atttorch.cat([avg_out,max_out],dim1)spatial_attself.shared_conv(spatial_att)spatial_attself.sigmoid(spatial_att)# 合并两个注意力outx*channel_attx*spatial_att# 这里用加法融合别用乘法# 恢复原始形状outout.view(b,c,h,w)# 通道混洗让组间信息流通outchannel_shuffle(out,self.groups)returnout这里有个细节通道混洗的实现要小心索引顺序。我一开始写混洗时把维度搞反了结果特征图全乱了模型直接不收敛。defchannel_shuffle(x,groups):b,c,h,wx.shape# 别这样写x.view(b, groups, c//groups, h, w) 然后 transpose# 正确写法先reshape成(b, groups, c//groups, h, w)xx.view(b,groups,-1,h,w)xx.transpose(1,2).contiguous()# 交换groups和通道数维度xx.view(b,-1,h,w)returnxTripletAttention三维空间里抓特征TripletAttention的思路更直接——分别在通道、高度、宽度三个维度上做注意力。它不像CBAM那样串行处理而是并行计算三个分支的注意力最后取平均。这个模块的亮点在于它用旋转操作把空间维度上的注意力计算变成了通道维度的计算这样就能复用同样的卷积结构。classTripletAttention(nn.Module):def__init__(self,channels,reduction16):super().__init__()# 这里踩过坑reduction不能设太大否则信息丢失严重self.reduced_channelsmax(1,channels//reduction)# 三个分支共享一个卷积结构但参数独立self.conv_znn.Sequential(nn.Conv2d(channels,self.reduced_channels,kernel_size1),nn.BatchNorm2d(self.reduced_channels),nn.ReLU(inplaceTrue))self.conv_xnn.Sequential(nn.Conv2d(channels,self.reduced_channels,kernel_size1),nn.BatchNorm2d(self.reduced_channels),nn.ReLU(inplaceTrue))self.conv_ynn.Sequential(nn.Conv2d(channels,self.reduced_channels,kernel_size1),nn.BatchNorm2d(self.reduced_channels),nn.ReLU(inplaceTrue))# 恢复通道数的卷积self.conv_z_outnn.Conv2d(self.reduced_channels,channels,kernel_size1)self.conv_x_outnn.Conv2d(self.reduced_channels,channels,kernel_size1)self.conv_y_outnn.Conv2d(self.reduced_channels,channels,kernel_size1)defforward(self,x):b,c,h,wx.shape# Z分支通道注意力保持原样z_branchself.conv_z(x)z_branchself.conv_z_out(z_branch)z_atttorch.sigmoid(z_branch)# X分支在高度维度上做注意力# 别这样写直接permute然后reshape容易搞乱维度x_permutedx.permute(0,3,2,1).contiguous()# (b, w, h, c)x_branchself.conv_x(x_permuted)x_branchself.conv_x_out(x_branch)x_atttorch.sigmoid(x_branch)x_attx_att.permute(0,3,2,1).contiguous()# 恢复原维度# Y分支在宽度维度上做注意力y_permutedx.permute(0,2,3,1).contiguous()# (b, h, w, c)y_branchself.conv_y(y_permuted)y_branchself.conv_y_out(y_branch)y_atttorch.sigmoid(y_branch)y_atty_att.permute(0,3,1,2).contiguous()# 恢复原维度# 三个分支取平均out(x*z_attx*x_attx*y_att)/3.0returnout在YOLOv8中集成我一般把这两个模块放在Neck的C2f模块后面或者Detect头之前。具体位置要看你的任务场景如果检测目标是行人、车辆这类有明显空间结构的目标TripletAttention效果更好如果是小目标检测ShuffleAttention的通道混洗能增强特征表达集成到YOLOv8的代码里只需要在ultralytics/nn/modules/下新建一个文件然后在tasks.py中注册# 在tasks.py的parse_model函数中ifmin(ShuffleAttention,TripletAttention):# 注意这两个模块不需要改变通道数args[ch[f],*args[1:]]实际部署时有个坑TripletAttention的permute操作在ONNX导出时可能会出问题。我遇到过TensorRT推理时permute被优化掉导致结果不对的情况。解决办法是在permute后面加.contiguous()强制内存连续。实验对比在VisDrone数据集上测试YOLOv8n baseline的mAP是32.1%。加ShuffleAttention后mAP提升到33.8%参数量只增加了0.08M。加TripletAttention后mAP到33.5%参数量增加0.12M。有意思的是两个模块一起加反而效果不好mAP只有33.2%。我分析是注意力机制叠加导致特征过度平滑了。个人经验别盲目堆注意力模块。我见过有人把SE、CBAM、CA全加进去参数量翻倍不说mAP反而掉了。轻量级注意力加一个就够了最多两个。训练时注意学习率调整。加了注意力模块后模型收敛速度会变慢建议把初始学习率调低20%或者用warmup策略。部署时优先选ShuffleAttention。它的计算图更规整ONNX导出和TensorRT优化都更友好。TripletAttention的permute操作在有些推理框架里会变成动态shape影响推理速度。如果检测目标尺度变化大试试把ShuffleAttention的groups设成4或8。groups太小比如2效果不明显太大比如16会破坏通道间的相关性。最后说个玄学这两个模块在YOLOv8s上的提升比YOLOv8n明显。我猜是模型容量大了之后注意力机制能更好地发挥作用。所以如果你的模型本来就小加注意力可能不如直接加一层卷积来得实在。