第一步,新建attention.py文件
attention.py放在那里可任意,本文放在ultralytics/nn/modules下面
之后在attention.py文件中添加各种注意力模块
第二步,修改yolov8.yaml文件
第一种是在yolov8.yaml文件中添加注意力机制层
(建议新建一个文件,文本为yolov8-EMA.yaml)
此时要注意,在加入注意力机制后,原先的层次会发现变化,在添加注意力机制层次之后的层次需要改动。
例如
原先的yolov8.yaml文件为
改动之后的yolov8.yaml文件为
第二种是改动yolov8.yaml文件中的已有的层次
首先,需要在block.py文件中,复制一份需要更改的层次(此处为C2f)然后改换一个名字(此处为C2f_Attention)
之后要在文件中增加注意力机制
注意,有些注意力机制需要表明输入的通道数
例如此处,c2为self.cv2的输出通道数,需要注入到SEAttention中
最后在yolov8.yaml文件修改层次
第三步,修改配置文件
对于第一种方法
#1,在ultralytics/nn/tasks.py文件中导入attention.py文件
from ultralytics.nn.modules.attention import EMA,SimAM
#2,在ultralytics/nn/tasks.py文件中找到def parse_model(d, ch, verbose=True):函数。在函数中添加如下的代码
elif m in(EMA,):
c2=ch[f]
args=[c2,*args]
如图
对于第二种方法
#1,在block.py文件中添加模块
#2,在ultralytics/nn/modules/__init__.py文件中添加模块
#3,添加模块在ultralytics/nn/tasks.py文件中
#4,在ultralytics/nn/tasks.py文件中找到def parse_model(d, ch, verbose=True):函数,添加模块
第四步,在程序中运行
第一种
model = YOLO("yolov8n-EMA.yaml")
第二种
model = YOLO("yolov8n.yaml")
最后运行结果为
第一种
第二种
可以发现注意力机制已经添加成功