保存模型加载分为两种:
1.只保存加载模型参数
torch.save(model.state_dict(), '\modelname.pkl')
# 加载到模型中,首先新建模型,然后加载参数
newmodel= NewModel(...)
newmodel.load_state_dict(torch.load('\modelname.pkl'))
2.保存完整模型
#保存模型
torch.save(model, '\modelname.pkl')
# 加载模型
model = torch.load('\modelname.pkl')