pytorch加载网络权重并冻结部分权重

需求

构建一个新的网络，加载一个已经训练好的网络的公共部分。

解决

import torch
from UnetModelsCuaseNotPadding_TwoStreamT2 import CRN_Net

def fix_bn(m):
    classname = m.__class__.__name__
    if classname.find('BatchNorm') != -1:
        m.eval()  
model = CRN_Net()
pretext_model = torch.load('../preTrain.pth.tar',"cpu").state_dict()
model_dict = model.state_dict()
print('new model',len(model_dict))
print('old model',len(pretext_model))        
#load part weight
state_dict = {k:v for k,v in pretext_model.items() if k in model_dict.keys()} #load same name layer weiget
dict_name = list(state_dict)
print('same pam',len(dict_name))
model_dict.update(state_dict)
model.load_state_dict(model_dict) 

#frozen part weight
for i,p in enumerate(model.parameters()):
    if i < 64:
        p.requires_grad = False 
        
# fix batchnorm
model.train()
model.apply(fix_bn) 

for name, p in model.named_parameters():
    print(f'{name}:\t{p.requires_grad}')```

值得注意的的是上面的64是加载公共部分的参数，其中不包括BN层的running_mean和running_var，所以自己需要认真和对下在确定。在训练的时候需要把BN层的所有参数都冻结上。

参考

pytorch那些坑——你确定你真的冻结了BN层