【修改 删除】修改神经网络权重key名称和删除权重

权重名称不匹配?权重参数shape不匹配?权重参数维度不匹配?

在这里插入图片描述
在这里插入图片描述

1 查看模型权重的key

    for k, v in model.state_dict().items():
        print(k)
        
    for k, v in weights_dict.items():
        print(k)

在这里插入图片描述

2 修改权重名称key

weights_dict_origin = torch.load(args.weights, map_location=device)['state_dict']

weights_dict = {
    
    
    k[len("backbone.") :] if "backbone." in k else k: v
    for k, v in weights_dict_origin.items()
}

3 删除某层权重参数

        # 删除有关分类类别的权重
        for k in list(weights_dict.keys()):
            if "output" in k:
                del weights_dict[k]
        print(model.load_state_dict(weights_dict, strict=False))

猜你喜欢

转载自blog.csdn.net/weixin_45377629/article/details/127451598