- 首先写好注意力机制网络结构,我以最简单的SE为例。
class SE(nn.Module):
def __init__(self,c1,c2,ratio=16) -> None:
super(SE,self).__init__()
self.avgpool = nn.AdaptiveAvgPool2d(1)
self.l1 = nn.Linear(c1,c1//ratio,bias=False)
self.relu = nn.ReLU(inplace=True)
self.l2 = nn.Linear(c1//ratio,c1,bias = False)
self.sig = nn.Sigmoid()
def forward(self,x):
b,c,_,_,=x.size()
y = self.avgpool(x).view(b,c)
y = self.l1(y)
y = self.relu(y)
y = self.l2(y)
y = self.sig(y)
y = y.view(b,c,1,1)
return x*y.expand_as(x)
由上图可以看出,SE注意力机制(squeeze and excitation networks) 经过一次平均池化,一次全连接,relu激活,一次全连接,sigmoid激活。个人感觉经过这一步可以对大目标或者小目标有更好的检测效果。
- 其次将这一网络添加到yolo中的common.py文件中
- 添加完后在models文件夹下新建一个.yaml文件,这里可以直接复制yolo5m.yaml然后在其backbone中的倒数第二层加入我们写好的SE。
- 最后一步,在yolo.py文件中找到图中所示位置,将定义好的SE网络添加到其中。
大功告成,接下来可以自己train一下试试水,我浅试了一下用coco128数据集在加了SE注意力机制后的效果,见下图。
由于轮数较少,所以效果不是很好,但可以明显看出回归框损失和loss损失有一丝丝效果,相信大数据训练效果会很明显。期待下次200epochs的效果展示。
- 最后的最后,其实可以加论文里各种注意力机制,只要按我上述方法操作。
- 除了上述添加方法,还可以在yolo配置文件里直接修改c3,但是操作有点繁琐且容易error,你们可以试试,成功了教教我。