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

最新推荐文章于 2025-02-08 03:28:35 发布

zheng shi wu ，

最新推荐文章于 2025-02-08 03:28:35 发布

阅读量6.7k

点赞数 8

文章标签：深度学习 python

本文链接：https://blog.csdn.net/weixin_54723793/article/details/119043506

版权

本文介绍了解决YoloV3网络模型结构修改后，使用原网络预训练权重出现不匹配的问题。通过将预训练权重视为有序字典，并采用for循环加载部分权重，再手动加载剩余权重的方法来解决这一问题。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

问题描述

修了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')