当前位置: 首页>编程语言>正文

python 下载预训练模型 pytorch预训练模型加载

1、 预训练模型网络结构 = 你要加载模型的网络结构
那么直接 套用

path="你的 .pt文件路径"
model = "你的网络"
checkpoint = torch.load(path, map_location=device)
model.load_state_dict(checkpoint)

2、 预训练模型网络结构 与你的网络结构不一致
当你直接套用上面公式,会出现类似unexpected key module.xxx.weight问题
这种情况下,需要具体分析一下网络信息,再决定如何加载。

# model_dict 是一个字典,保存网络 各层名称和参数,
model_dict = model.state_dict()
print(model_dict.keys()
# 这里打印出 网络 各层名称
checkpoint = torch.load(path,map_location=device)
for k, v in checkpoint.items():
    print("keys:".k)
# 这里打印出 预训练模型网络 各层名称, 是字典 【键】显示的另一种方式。

然后,对比两者网络结构参数 的异同,

若各层网络名称 基本不一致,那这个预训练模型基本就没法用了,直接换模型吧
若两者网络参数有很多 类似的地方,但又不完全一致,那可以采取如下方式。
(1) 部分网络关键字 ---- 完全匹配的情况

model.load_state_dict(checkpoint, strict=True)

load_state_dict 函数添加 参数 strict=True, 它直接忽略那些没有的dict,有相同的就复制,没有就直接放弃赋值!他要求预训练模型的关键字必须确切地严格地和 网络的 state_dict() 函数返回的关键字相匹配才能赋值。
strict 也不是很智能,适用于那些 网络关键字 基本能够匹配的情况。否则即使加载成功,网络参数也是空的。
(2)大部分网络关键字 ---- 部分匹配 (不完全相同,但类似),例如
网络关键字: backbone.stage0.rbr_dense.conv.weight
预训练模型 关键字:stage0.rbr_dense.conv.weight
可以看到,网络关键字 比预训练模型 多了一个前缀,其它完全一致,这种情况下,可以把 预训练模型的 stage0.rbr_dense.conv.weight 读入 网络的 backbone.stage0.rbr_dense.conv.weight 中。

# 对于 字典而言,in 或 not in 运算符都是基于 key 来判断的
model_dict = model.state_dict()
checkpoint = torch.load(path,map_location=device)
# k 是预训练模型的一个关键字, ss是 网络的有一个关键字
for k, v in checkpoint.items():
    flag = False
    for ss in  model_dict.keys():
        if k in ss:  # 在每一个元素内部匹配
            s = ss; flag = True; break
        else:
            continue
    if flag:
        checkpoint[k] = model_dict[s]

3、断点恢复
我感觉这个和常规【模型保存加载】方法的区别主要是 epoch的恢复

# 模型保存
state = {
    'epoch': epoch,
    'state_dict': model.state_dict(),
    'optimizer': optimizer.state_dict(),
     ... # 有其他希望保存的内容,也可自定义
    }
    torch.save(state, filepath)
# 加载模型,恢复训练
    model.load_state_dict(state['state_dict'])
    optimizer.load_state_dict(state['optimizer'])
       start_epoch = checkpoint['epoch'] + 1

https://www.xamrdz.com/lan/5n31925002.html

相关文章: