pytorch模型的保存及加载

pytorch模型的保存及加载主要用到下面几行代码

#模型的保存,在模型训练完以后

torch.save(model.state_dict(),'./model.pt')

#模型的加载
model1=simpleNet(28*28,10,300,100)#输入原本模型所需的参数，定义的网络的名称为simpleNet,如果是在同一个程序中运行，直接调用就可以，如果不是同一个程序中，那么需要在该语句前重新定义一遍simpleNet网络。

model1.load_state_dict(torch.load('./model.pt'),strict=True)

model1.eval()