pytorch 保存模型和恢复模型

主要有两种方法序列化和恢复模型

(1)(推荐)只保存和加载模型参数

torch.save(the_model.state_dict(), PATH)
#恢复
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH)

(2)保存和加载整个模型(然而,在这种情况下,序列化的数据被绑定到特定的类和固定的目录结构,所以当在其他项目中使用时,或者在一些严重的重构器之后它可能会以各种方式break )

torch.save(the_model, PATH)
#恢复
the_model = torch.load(PATH)

猜你喜欢

转载自www.cnblogs.com/come-on-baby/p/11042019.html