yolo模型修改后怎么加载预训练权重

问题描述

• 修了yolov3网络模型结构，使用原网络的预训练权重会出现 mismatch 的情况

解决办法

• 总体思想就是把预训练权重当成一个有序字典处理，其实本质上它就是一个有序字典。有序字典的定义及用法这里不赘述，不了解的自己去查。
• 采用for循环，加载部分权重。
• 手动加载剩余权重

代码实现

# 载入与训练权重，此时并未载入网络
ckpt = torch.load(weights, map_location=device)
# 取出网络的模型，其中里面的参数是调用:_initialize_weights(self)生成的
model_state_dict = model.state_dict()
# 将前87层预训练权重赋予模型参数
 for i, (k, v) in enumerate(ckpt["model"].items()):
    if i < 348:
        model_state_dict[k] = v
# 手动添加不匹配参数
model_state_dict['module_list.90.Conv2d.weight'] = ckpt["model"]['module_list.87.Conv2d.weight']
model_state_dict['module_list.90.BatchNorm2d.weight'] = ckpt["model"]['module_list.87.BatchNorm2d.weight']
model_state_dict['module_list.90.BatchNorm2d.bias'] = ckpt["model"]['module_list.87.BatchNorm2d.bias']
... ...
# 保存权重
torch.save(model_state_dict, 'weights.pt')

---

对于结构删减改动较大的网络，手动加载也挺麻烦的，若找到新的方法，会继续更新