1.加载全部模型:
net.load_state_dict(torch.load(net_para_pth))
2.加载部分模型
net_para_pth = './result/5826.pth'
pretrained_dict = torch.load(net_para_pth)
model_dict = net.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
net.load_state_dict(model_dict)
3.改变某一层参数后加载
将该层名称改一下,然后用2中方法加载,比如,要将conv5的out_channels由256改为512。
将conv_5改为conv_5_chg,就可以顺利加载了,不改会报错哟
算是权宜之计了,还有什么好方法,希望多多指教
4.单GPU/CPU加载多GPU训练的网络
正常情况下,多GPU保存模型应该加上.module,然后加载时即使是单GPU也不会出问题,但是如果保存时忘记加,加载时就需要多一道手续
参考:javascript:void(0)
def load_GPU(model, model_path, mapLoc='cpu'):
state_dict = torch.load(model_path, map_location=mapLoc)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
return model
也就是重新定义一个OrderedDict,然后将state_dict键值中的.module去掉