保存和加载模型的时候,配对的函数: 对于仅保存state_dict()的方式,那保存和加载模型的方式为: 保存:torch.save(model.state_dict(), PATH) 加载:model.laod_state_dict(torch.load(PATH)) 一般加载模型是在训练完成后用模型做测试,这时候加载模型记得要加上model.eval(),把模型切换到evaluation模式,这时候会调整dropout和bactch的模式。 对于保存和加载整个模型的情况: torch.save(model, PATH) model = torch.load(PATH) 可以看到,前面的model.load_state_dict()和这里的不同,前面的情况需要你先定义一个模型,然后再load_state_dict() 但是这里load整个模型,会把模型的定义一起load进来。完成了模型的定义和加载参数的两个过程。
pytorch 保存和加载模型
最新推荐文章于 2023-07-28 10:37:26 发布