YOLO添加注意力机制

这里介绍下如何为YOLO模型添加注意力机制,主要在于卷积神经网络CNN劣势在于无法很好关注重要的位置。而通过引入注意力机制,从而加强其对一些物品的注意。

之前的文章已经介绍过YOLO模型的结构,要为模型添加注意力机制,可以从如下3个方面下手:

下面从第一步开始介绍。相应的注意力机制模块主要有:

SENet,让网络自己学会哪些通道重要,对于不重要的通道进行抑制
CBAM,可以看成通道注意力+空间注意力的集成,在上一步基础上添加看图中哪些位置的能力
BAM,瓶颈注意力模块,与CBAM不同,其使用并联结构
Non-Local Attention或Non-Local Neural Network,是一种全局建模方法,通过计算任意两个位置之间的关系,实现信息的长距离传递
Transformer风格注意力,如Contextual Transformer,兼顾局部细节提取和上下文建模。而ECA(Efficient Channel Attention)是对SENet的改进,计算更快。而Coordinate Attention更适用于移动设备和轻量网络

而引入注意力模块,可以在ultralytics/nn目录下新建对应的模块,并编写对应网络结构代码。而在ultralytics/nn/tasks.py中导入之前的模块,修改parse_model函数中内容,添加相应加载方式。

接着为了避免与之前模型结构冲突,在ultralytics/cfg/models/对应YOLO版本目录下新建1个yaml文件,其中backbone中添加相应的注意力机制模块,并修改head中网络的索引。

一切准确就绪就可以加载模型并调用其train方法进行训练了,此时配置文件设置为之前新建的模型yaml文件。

可以看到整个流程还是比较简单,而可视化注意力可以考虑使用visualiz包,通过以下方式进行安装:

pip install visualize==0.5.0

参考文章:

https://blog.csdn.net/shuai_beibei/article/details/142798188

https://blog.csdn.net/2403_88150975/article/details/148263735

https://www.cnblogs.com/Brisling/p/16838987.html

https://blog.csdn.net/weixin_52603404/article/details/145930505

https://zhuanlan.zhihu.com/p/356798637

如果喜欢这篇文章或对您有帮助，可以：[☕] 请我喝杯咖啡 | [💓] 小额赞助

码力全开 / YOLO添加注意力机制