当前位置 博文首页 > pytorch 实现冻结部分参数训练另一部分
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)
查看每一层的代码:
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 gamma 1 resnet.conv1.weight 2 resnet.bn1.weight 3 resnet.bn1.bias 4 resnet.bn1.running_mean 5 resnet.bn1.running_var 6 resnet.layer1.0.conv1.weight 7 resnet.layer1.0.bn1.weight 8 resnet.layer1.0.bn1.bias 9 resnet.layer1.0.bn1.running_mean ....
同样在模型中添加这样的代码：
for i,p in enumerate(net.parameters()): if i < 165: p.requires_grad = False
在优化器中添加上面的那句话可以实现参数的屏蔽
补充：pytorch 加载预训练模型 + 断点恢复 + 冻结训练（避坑版本）
那么直接 套用
path="你的 .pt文件路径" model = "你的网络" checkpoint = torch.load(path, map_location=device) model.load_state_dict(checkpoint)
当你直接套用上面公式，会出现类似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)
# 这里打印出 预训练模型网络 各层名称, 是字典 【键】显示的另一种方式。
然后，对比两者网络结构参数 的异同，
若各层网络名称 基本不一致，那这个预训练模型基本就没法用了，直接换模型吧
若两者网络参数有很多 类似的地方，但又不完全一致，那可以采取如下方式。
model.load_state_dict(checkpoint, strict=True)
load_state_dict 函数添加 参数 strict=True， 它直接忽略那些没有的dict，有相同的就复制，没有就直接放弃赋值！他要求预训练模型的关键字必须确切地严格地和 网络的 state_dict() 函数返回的关键字相匹配才能赋值。
strict 也不是很智能，适用于那些 网络关键字 基本能够匹配的情况。否则即使加载成功，网络参数也是空的。
网络关键字： 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]
我感觉这个和常规【模型保存加载】方法的区别主要是 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
一般冻结训练都是针对【backbone】来说的，较多应用于【迁移学习】
例如，0-49 Epoch：冻结 backbone进行训练；50-99：不冻结训练。
Init_Epoch = 0 Freeze_Epoch = 50 Unfreeze_Epoch =100 #------------------------------------# # 冻结一定部分训练 #------------------------------------# for param in model.backbone.parameters(): param.requires_grad = False for epoch in range(Init_Epoch,Freeze_Epoch): # I`m Freeze-training !! pass #------------------------------------# # 解冻后训练 #------------------------------------# for param in model.backbone.parameters(): param.requires_grad = True for epoch in range(Freeze_Epoch,Unfreeze_Epoch): # I`m unfreeze-training !! pass
以上为个人经验，希望能给大家一个参考，也希望大家多多支持站长博客。如有错误或未考虑完全的地方，望不吝赐教。
js