pytorch 实现冻结部分参数训练另一部分

时间:2021-05-23

1)添加下面一句话到模型中

for p in self.parameters(): p.requires_grad = False

比如加载了resnet预训练模型之后,在resenet的基础上连接了新的模快,resenet模块那部分可以先暂时冻结不更新,只更新其他部分的参数,那么可以在下面加入上面那句话

class RESNET_MF(nn.Module): def __init__(self, model, pretrained): super(RESNET_MF, self).__init__() self.resnet = model(pretrained) for p in self.parameters(): p.requires_grad = False #预训练模型加载进来后全部设置为不更新参数,然后再后面加层 self.f = SpectralNorm(nn.Conv2d(2048, 512, 1)) self.g = SpectralNorm(nn.Conv2d(2048, 512, 1)) self.h = SpectralNorm(nn.Conv2d(2048, 2048, 1)) ...

同时在优化器中添加:

filter(lambda p: p.requires_grad, model.parameters())optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001, \ betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-5)

2) 参数保存在有序的字典中,那么可以通过查找参数的名字对应的id值,进行冻结

查看每一层的代码:

model_dict = torch.load('net.pth.tar').state_dict()dict_name = list(model_dict)for i, p in enumerate(dict_name): print(i, p)

打印一下这个文件,可以看到大致是这个样子的:

0 gamma1 resnet.conv1.weight2 resnet.bn1.weight3 resnet.bn1.bias4 resnet.bn1.running_mean5 resnet.bn1.running_var6 resnet.layer1.0.conv1.weight7 resnet.layer1.0.bn1.weight8 resnet.layer1.0.bn1.bias9 resnet.layer1.0.bn1.running_mean....

同样在模型中添加这样的代码:

for i,p in enumerate(net.parameters()): if i < 165: p.requires_grad = False

在优化器中添加上面的那句话可以实现参数的屏蔽

补充: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

4、冻结训练

一般冻结训练都是针对【backbone】来说的,较多应用于【迁移学习】

例如,0-49 Epoch:冻结 backbone进行训练;50-99:不冻结训练。

Init_Epoch = 0Freeze_Epoch = 50Unfreeze_Epoch =100#------------------------------------## 冻结一定部分训练#------------------------------------# for param in model.backbone.parameters(): param.requires_grad = Falsefor epoch in range(Init_Epoch,Freeze_Epoch): # I`m Freeze-training !! pass#------------------------------------## 解冻后训练#------------------------------------#for param in model.backbone.parameters(): param.requires_grad = Truefor epoch in range(Freeze_Epoch,Unfreeze_Epoch): # I`m unfreeze-training !! pass

以上为个人经验,希望能给大家一个参考,也希望大家多多支持。如有错误或未考虑完全的地方,望不吝赐教。

声明:本页内容来源网络,仅供用户参考;我单位不保证亦不表示资料全面及准确无误,也不保证亦不表示这些资料为最新信息,如因任何原因,本网内容或者用户因倚赖本网内容造成任何损失或损害,我单位将不会负任何法律责任。如涉及版权问题,请提交至online#300.cn邮箱联系删除。

相关文章