修改 YOLOv5 7.0 的backbone—以ResNet为例

引言

借鉴魔傀面具老师和霹雳吧啦Wz老师的代码和课程讲解，完成了不使用Timm库的backbone修改。

两位老师的bilibili主页和github主页分别如下：

bilibili：

        魔傀面具的个人空间-魔傀面具个人主页-哔哩哔哩视频

        霹雳吧啦Wz的个人空间-霹雳吧啦Wz个人主页-哔哩哔哩视频

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方式。分别引自博主

赢勾喜欢海 和 落难Coder

下面，开始逐步修改。

一、准备模型文件和配置文件

直接从霹雳吧啦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上都有：

https://github.com/z1069614715/objectdetection_script/blob/master/yolo-improve/yolov5-backbone/yolo.py

导入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文件：

https://github.com/z1069614715/objectdetection_script/blob/master/yolo-improve/yolov5-backbone/yolov5-custom.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部分的预训练权重。