添加注意力机制
各种注意力模块优缺点
优点: 可以通过学习自适应的通道权重,使得模型更加关注有用的通道信息。 缺点: SE注意力机制只考虑了通道维度上的注意力,无法捕捉空间维度上的注意力,适用于通道数较多的场景,但对于通道数较少的情况可能不如其他注意力机制。
在模型中添加注意力模块
设置备选注意力模块,在nets/yolo.py里添加如下代码:
添加加载模型时所需的选取注意力的参数,在nets/yolo.py里YoloBody类添加phi_attention参数:(例如本例中,phi_attention为0代表不使用注意力机制,1-4代表上述备选列表中的4个模块)
在模型中添加注意力模块(本例所示为在P3层添加注意力模块,类型通过phi_attention参数指定,您可自行根据需要修改):
在__init__函数中进行注意力模块的初始化
在forward中选择需要插入注意力的位置,例如在特征金字塔P3层后添加注意力机制
使用注意力模块改进的模型
在train.py中模型创建部分修改phi_attention参数,按照上节所述步骤训练模型。
在yolo.py中修改phi_attention参数,加载训练后的权重即可正常使用改进后的模型。
Last updated