本文记录自己修改yolov8模型结构的过程,主要展示在yolov8模型结构中添加CBAM注意力模块的方式,另一个原因是CBAM注意力模块在实际应用也确实有效果;后续其他修改可以根据需要按照以下类似方式进行模块添加和修改即可。
yolov8模型结构
在对yolov8模型结构进行修改前,首先,需要熟悉yolov8整体的网络结构,在此基础之上,根据需要在特定位置添加自定义的网络模块才不会出错。
yolov8模型结构简要说明:
1)根据下方的架构图,将yolov8分为3个组成部分:backbone、neck和head;
2)backbone:表示骨干网络
,负责从输入图像中提取特征,将图像转化为具有丰富语义信息的特征表示;
3)neck:表示中间网络
,用于对来自 backbone 的特征进行多尺度特征融合和增强,以提升模型的性能;
4)head:表示目标检测头,使用边界框回归器进行边界框回归和分类器进行目标分类。
yolov8模型架构(图片来自MMYOLO)
修改模型结构
以下以添加CBAM注意力模块为例,并且yolov8代码中也集成了CBAM模块的实现,以此为参考记录整个修改过程。
1、添加自定义模块
1)在yolov8中添加CBAM注意力模块,首先打开脚本ultralytics/nn/modules/conv.py
,脚本路径如下所示:
** 注意:自定义的网络模块不是固定放到conv.py脚本中,我们也可以自定义新的脚本添加模块,但是要按下面的方式引用和声明。**
2)并在脚本中添加CBAM模块的模型代码,如下所示:
3)同时在脚本__all__变量中添加可导出的类名,如下所示:
2、添加自定义模块声明
添加CBAM模块的代码后,还需要在ultralytics/nn/modules/__init__.py
脚本中添加自定义模块声明,用于后续引用,如下所示:
脚本路径:
在ultralytics/nn/modules/__init__.py
脚本中添加模块声明:
3、解析模型函数逻辑修改
1)解析模型函数解读
yolov8中解析模型函数用于搭建网络模型,解析模型函数对应的路径为ultralytics/nn/tasks.py
。首先解读解析模型函数,通过熟悉该函数功能逻辑,帮助我们更好的添加自定义模块逻辑。具体解读代码注释如下:
def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)
"""
函数功能:更新当前层的args,计算当前层的输出通道c2,使用当前层的参数构建当前层
入参说明:
d (dict): 解析xx.yaml模型配置文件的字典形式
ch int: 输入的channel数,后续的ch是记录模型每一层的输出通道数
出参说明:
nn.Sequential(*layers): 网络的每一层的层结构
sorted(save): 把所有层结构中from不是-1的值保存下来并根据序号排序
"""
import ast
# Args
max_channels = float("inf")
nc, act, scales = (d.get(x) for x in ("nc", "activation", "scales"))
depth, width, kpt_shape = (d.get(x, 1.0) for x in ("depth_multiple", "width_multiple", "kpt_shape"))
# 获取模型规模信息:模型深度数、宽度数、最大通道数,
# 对应5种规模:n/s/m/l/x
if scales:
scale = d.get("scale")
if not scale:
scale = tuple(scales.keys())[0]
LOGGER.warning(f"WARNING ⚠️ no model scale passed. Assuming scale='{
scale}'.")
depth, width, max_channels = scales[scale]
# 设置激活函数
if act:
Conv.default_act = eval(act) # redefine default activation, i.e. Conv.default_act = nn.SiLU()
if verbose:
LOGGER.info(f"{
colorstr('activation:')} {
act}") # print
if verbose:
LOGGER.info(f"\n{
'':>3}{
'from':>20}{
'n':>3}{
'params':>10} {
'module':<45}{
'arguments':<30}")
# 开始搭建网络
# ch: 记录模型每一层的输出通道数
# layers: 保存每一层的层结构
# save: 记录所有层结构中from中不是-1的层结构序号
# c2: 保存当前层的输出channel
ch = [ch]
layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out
# d["backbone"] + d["head"],将配置文件中backbone和head结构信息组合在一起
# from:表示当前层输入来自哪几层
# number:表示当前层数
# module:表示当前层名称
# args:表示当前入参参数
for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]): # from, number, module, args
# 从torch.nn获取对应名称的网络模块
m