pytorch model save and load

There are two types of
saving and loading models : 1. Only save and load model parameters

torch.save(model.state_dict(), '\modelname.pkl')
# 加载到模型中,首先新建模型,然后加载参数
newmodel= NewModel(...)
newmodel.load_state_dict(torch.load('\modelname.pkl'))

2. Save the complete model

#保存模型
torch.save(model, '\modelname.pkl')
# 加载模型
model = torch.load('\modelname.pkl')

Guess you like

Origin blog.csdn.net/qq_43360777/article/details/106311178