# (本例子情况:A中的(B)的权重key比B的结构名中多了'res.' 可类比修改)
# 模型A中存在的结构B权重文件路径:
weight = '权重路径'
# 模型B:以res50为例
model = torchvision.models.resnet50(num_classes).to(device)
# 建一个空字典,模拟权重文化
state_dict = dict()
# 把模型A的权重中 key-value 的 key 打印出来
print('[pretrained_model_keys]:', *torch.load(weight).keys(), sep="\n")
for k, v in torch.load(weight).items():
# 将多余出来的res. 替换为空
state_dict[k.replace('res.', '')] = v
# 加载即可
# model.load_state_dict(torch.load(weight), strict=False)
# missing_keys代表该模型中没有匹配到权重的层
# unexpected_keys代表 所加载的权重文件中,未匹配到层结构的权重名
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
print('[missing_keys]:', *missing_keys, sep="\n")
print('[unexpected_keys]:', *unexpected_keys, sep="\n")
模型B加载模型A部分结构(B)的权重参数,保证机构名与权重名对齐加载方法
最新推荐文章于 2024-05-06 10:23:59 发布