引言
借鉴魔傀面具老师和霹雳吧啦Wz老师的代码和课程讲解,完成了不使用Timm库的backbone修改。
两位老师的bilibili主页和github主页分别如下:
bilibili:
github:
魔傀老师@ https://github.com/z1069614715
霹雳吧啦Wz老师@ https://github.com/WZMIAOMIAO
CSDN:
霹雳吧啦Wz老师@ 太阳花的小绿豆
写下来主要是防止自己忘记,也可以供大家参考,如果有大佬能提些建议是更好啦。
简介
基于Ultralytics版的YOLOv5 7.0,修改YOLOv5的backbone有两种方式:
①对于常见的分类网络,如:VggNet、ResNet、ViT、SwinTransformer等。可以直接利用Timm库修改;
②对于一些不太常见的分类网络,则需要手动修改模型文件内容,利用yolo.py的“parse_model”函数建立修改后的模型。
对于上述第一种方式,直接看魔傀老师的讲解视频即可:
YOLOV5改进-基于TIMM更换你想要的主干网络(基本支持现有大部分CNN网络)_哔哩哔哩_bilibili
其实Timm库已经包含了ResNet分类网络,直接按照上述讲解视频就可以运行成功。
但是,有一种情况,是针对特殊的检测任务,比如我在做的缺陷分类需要将图像分辨率初始化为640x640(或者更大的尺度),而Timm库给的权重应该是224x224的,这样一来就造成预训练权重和网络输入图像尺度不一致的问题。
同时呢,魔傀老师的github上没有ResNet的配置文件,我也看到了一些在配置文件中逐个添加Block的方法,如图1所示。可是我总觉得这种方式不太方便,所以就借鉴魔傀老师利用FasterNet网络修改backbone的方法,实现了一步式修改backbone。
图1. 修改YOLOv5的backbone方式。分别引自博主
下面,开始逐步修改。
一、准备模型文件和配置文件
直接从霹雳吧啦Wz老师的github主页复制resnet的model文件,粘贴到YOLOv5的models文件夹下。
从魔傀老师的github主页找到fasternet文件,用作修改的参考物。
图2. resnet模型文件。引自博主
图3. fasternet文件,包含fasternet.py和文件夹fasternet_cfg。引自up主 魔傀面具
出于方便,将resnet的model文件命名为:resnet.py。并仿照fasternet_cfg,新建一个resnet_cfg文件夹,里面存放resnet网络的参数配置,如图4所示。
图4. 所有新建的文件
接下来,详细讲解resnet_cfg配置文件夹的单个文件内有哪些内容。
首先参考 fasternet_t0.yaml 和 fastnet.py ,搞清楚yaml里到底写的是什么,两者对应关系如图5所示。
图5. 配置文件的参数关系
可以看出,yaml里写的是模型的初始化参数,所以我们有样学样,查看resnet网络需要的初始化参数。
图6. resnet模型的初始化参数
然后在新建的resnet_cfg文件夹里建立resnet34、resnet50、resnet101的参数文件,resnet34网络的内容如下图所示。resnet50和resnet101网络把blocks_num修改就行。
include_top指是否需要分类层,因为我们只是将resnet用来提取特征,所以这里设置为False。
图7. 写入resnet初始化参数
需要注意的是,参数文件中的block不能直接写“BasicBlock”或“Bottleneck”。因为,根据resnet的_make_layer函数,传入的参数block得是一个类,可分为:
Bottleneck、 BasicBlock
图8. block块的具体属性
如果直接输入“BasicBlock”或“Bottleneck”,那么会出现以下报错:
AttributeError: 'str' object has no attribute 'expansion'
二、修改模型文件
需要修改的文件有:
resnet.py(resnet的model.py)
yolo.py
yolov5s.yaml(也可以是n、m、l等变体)
借鉴的文件是:
fasternet.py
2.1 修改resnet.py
首先要知道的是,YOLOv5有4个特征层:P2、P3、P4、P5。如果简单地载入backbone的话,那就会缺少这4个特征层,后续的特征融合和检测头就无法发挥作用。
所以,需要从修改的backbone中提取出4个特征层,用作P2、P3、P4、P5。
以fasternet网络结构为例,我们来看看魔傀老师是怎样做的:
图9. fasternet网络结构
可以看到,fasternet分别有4个stage,可以分别对应于P2、P3、P4、P5。接下来,再在fasternet.py中看看魔傀老师的代码。
图10. 分阶段保存特征输出
魔傀老师太强辣!
于是,依据resnet网络结构,在resnet.py文件中也添加类似的代码。
图11. resnet网络结构
在resnet.py中,修改class Resnet的前向函数:
def forward(self, x) -> Tensor:
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
outs = []
for idx, layer in enumerate(self.layers):
x = layer(x)
outs.append(x)
return outs
并在末尾添加构建模型的调用函数:
def resnet34_(weights=None, cfg='resnet_cfg/resnet34.yaml'): # channel -> [64, 128, 256, 512]
with open(cfg) as f:
cfg = yaml.load(f, Loader=yaml.SafeLoader)
model = ResNet(block=BasicBlock, **cfg)
if weights is not None:
pretrain_weight = torch.load(weights, map_location='cpu')
model.load_state_dict(update_weight(model.state_dict(), pretrain_weight))
return model
def resnet50_(weights=None, cfg='resnet_cfg/resnet50.yaml'): # channel -> [256, 512, 1024, 2048]
with open(cfg) as f:
cfg = yaml.load(f, Loader=yaml.SafeLoader)
model = ResNet(block=Bottleneck, **cfg)
if weights is not None:
pretrain_weight = torch.load(weights, map_location='cpu')
model.load_state_dict(update_weight(model.state_dict(), pretrain_weight))
return model
def resnet101_(weights=None, cfg='models/resnet_cfg/resnet101.yaml'): # channel -> [256, 512, 1024, 2048]
with open(cfg) as f:
cfg = yaml.load(f, Loader=yaml.SafeLoader)
model = ResNet(block=Bottleneck, **cfg)
if weights is not None:
pretrain_weight = torch.load(weights, map_location='cpu')
pretrain_weight = pretrain_weight['state_dict']
model.load_state_dict(update_weight(model.state_dict(), pretrain_weight))
return model
2.2 修改yolo.py
修改parse_model函数,以及_forward_once函数,这两个函数的代码在魔傀老师的github上都有:
导入resnet模块,然后在parse_model函数中找到类似位置,添加这样一段代码:
from models.resnet import * # 导入resnet模块
elif m in {resnet34_, resnet50_, resnet101_}:
m = m(*args)
c2 = m.channel
2.3 修改yolov5s.yaml
参考魔傀老师的yolov5.yaml文件:
简单修改1个地方就行,args为空。
图11. 修改yolov5s.yaml文件
三、效果展示
yolo.py中,打开
profile
line-profile
运行,即可输出网络结构和推理时间。
图12. 修改backbone后的YOLOv5s网络结构
图12. 修改backbone后的YOLOv5s网络推理时间和运算量
从图11和图12可以看到,已经成功将YOLOv5s的backbone修改为ResNet101。
此外,我也用Timm库修改过,进行对比后发现,Timm库修改后的网络推理时间会快一些,而参数量和运算量是保持一致的。
四、加载预训练权重
目前在用Timm库修改后的YOLOv5s进行训练,预训练权重在224x224图像上训练得到,并通过timm的接口载入,而我的模型输入初始化是640x640,检测精度很不理想hhh,现在还没搞清楚怎么回事。
对于我这种修改方法,我参考fasternet.py和yolov5的权重载入方法,测试了一下,是能够载入的。
权重载入代码如下所示,添加在resnet.py中。
def update_weight(model_dict, weight_dict):
idx, temp_dict = 0, {}
for k, v in weight_dict.items():
if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
temp_dict[k] = v
idx += 1
model_dict.update(temp_dict)
print(f'loading weights... {idx}/{len(model_dict)} items')
return model_dict
def intersect_dicts(da, db, exclude=()):
# Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
return {k: v for k, v in da.items() if k in db and all(x not in k for x in exclude) and v.shape == db[k].shape}
if __name__ == '__main__':
weights = '../resnet101.pth'
model = resnet101_(weights=None, cfg='resnet_cfg/resnet101.yaml')
print(model.channel)
if weights is not None:
pretrain_weight = torch.load(weights, map_location='cpu')
# state = pretrain_weight['model'].float().state_dict()
state = pretrain_weight['state_dict']
csd = intersect_dicts(state, model.state_dict())
model.load_state_dict(update_weight(model.state_dict(), pretrain_weight['state_dict']))
inputs = torch.randn((1, 3, 640, 640))
for i in model(inputs):
print(i.size())
结果输出:
loading weights... 624/624 items
正式载入权重:
上述测试成功后,为了能在YOLOv5的train.py中载入权重,直接在模型配置的yaml文件中写入权重路径即可。
图13. 载入backbone的权重
权重文件跟train.py脚本在同一目录下。
载入权重后,会在命令行中提示载入了多少权重,如下图所示。
图14. 权重载入的数量比
五、结论
通过载入backbone网络文件和修改YOLOv5的模型构建函数,成功用ResNet网络替换YOLOv5的backbone,并将ResNet网络的各个阶段输出分为4个特征层,符合YOLOv5网络整体结构。相较于Timm库的一键式操作,这种分步添加的方式更有利于局部模块的修改。但除了修改backbone之外,还需要考虑预训练权重的载入方式,包括neck和head部分的预训练权重。