pytorch 并行训练之后的模型

报错:Unexpected key(s) in state_dict: “module.conv1.0.weight”
这是由于并行训练导致的
解决办法,读取模型的时候,去掉module的前缀就好了

from collections import OrderedDict
state_dict=torch.load(’./trained_model/dense169_110.pth’)
new_state_dict = OrderedDict()
for k, v in state_dict.items():
if k ==“g”:
new_state_dict[k] = v

name = k[7:] # remove `module.`
new_state_dict[name] = v

torch.save(new_state_dict, ‘model_out.pth’)
model.load_state_dict(new_state_dict)#加载训练好的模型文件

你可能感兴趣的:(pytorch 并行训练之后的模型)