前 言:作为当前先进的深度学习目标检测算法YOLOv5,已经集合了大量的trick,但是在处理一些复杂背景问题的时候,还是容易出现错漏检的问题。此后的系列文章,将重点对YOLOv5的如何改进进行详细的介绍,目的是为了给那些搞科研的同学需要创新点或者搞工程项目的朋友需要达到更好的效果提供自己的微薄帮助和参考。
解决问题:加入Coordinate注意力机制,可以让网络更加关注待检测目标,提高检测效果
添加方法:
第一步:确定添加的位置,作为即插即用的注意力模块,可以添加到YOLOv5网络中的任何地方。
第二步:common.py构建CoTAttention模块。
class CoordAttn(nn.Module):
def __init__(self, dim=512,kernel_size=3):
super().__init__()
self.dim=dim
self.kernel_size=kernel_size
self.key_embed=nn.Sequential(
nn.Conv2d(dim,dim,kernel_size=kernel_size,padding=kernel_size//2,groups=4,bias=False),
nn.BatchNorm2d(dim),
nn.ReLU()
)
self.value_embed=nn.Sequential(
nn.Conv2d(dim,dim,1,bias=False),
nn.BatchNorm2d(dim)
)
factor=4
self.attention_embed=nn.Sequential(
nn.Conv2d(2*dim,2*dim//factor,1,bias=False),
nn.BatchNorm2d(2*dim//factor),
nn.ReLU(),
nn.Conv2d(2*dim//factor,kernel_size*kernel_size*dim,1)
)
def forward(self, x):
bs,c,h,w=x.shape
k1=self.key_embed(x) #bs,c,h,w
v=self.value_embed(x).view(bs,c,-1) #bs,c,h,w
y=torch.cat([k1,x],dim=1) #bs,2c,h,w
att=self.attention_embed(y) #bs,c*k*k,h,w
att=att.reshape(bs,c,self.kernel_size*self.kernel_size,h,w)
att=att.mean(2,keepdim=False).view(bs,c,-1) #bs,c,h*w
k2=F.softmax(att,dim=-1)*v
k2=k2.view(bs,c,h,w)
return k1+k2
第三步:yolo.py中注册CoordAtt模块
elif m is CoordAtt:
inp = args[0]
inp = make_divisible(inp * gw, 8) if inp != no else inp
args = [inp]
第四步:修改yaml文件,本文以修改head(特征融合网络)为例,将原C3模块后加入该模块。
第五步:将train.py中改为本文的yaml文件即可,开始训练。
结 果:本人在多个数据集上做了大量实验,针对不同的数据集效果不同,同一个数据集的不同添加位置方法也是有差异,需要大家进行实验。有效果有提升的情况占大多数。
预告一下:下一篇内容分享ECA轻量级通道注意力机制。有兴趣的朋友可以关注一下我,有问题可以留言或者私聊我哦
PS:Coordinate注意力机制,不仅仅是可以添加进YOLOv5,也可以添加进任何其他的深度学习网络,不管是分类还是检测还是分割,主要是计算机视觉领域,都可能会有不同程度的提升效果。
最后,希望能互粉一下,做个朋友,一起学习交流。