1. 预训练参数保存
def savecheckpoint(model,dir):
para=model.state_dict()
torch.save(para, dir)
2. 预训练参数加载
def LoadParameter(_structure, _parameterDir):
"""
:param _structure: model
:param _parameterDir: 参数位置
:return:
"""
checkpoint = torch.load(_parameterDir)
pretrained_state_dict = checkpoint['state_dict']
model_state_dict = _structure.state_dict()
for key in pretrained_state_dict:
if key in model_state_dict:
model_state_dict[key] = pretrained_state_dict[key]
"""
another method:
pretrained_dict = {k: v for k, v in pretrained_state_dict.items() if k in model_state_dict}
model_state_dict.update(pretrained_dict)
"""
_structure.load_state_dict(model_state_dict)
model = torch.nn.DataParallel(_structure).cuda()
return model