Pytorch加载和保存模型

Posted 2022-10-12 yqpy

首先注明，参考了这篇博客https://www.jianshu.com/p/4905bf8e06e5

• 方法1、

先序列化，格式可以是mdl，pt等

torch.save(model.state_dict(), MODEL_PATH)

然后反序列化，再加载

model.load_state_dict(torch.load(MODEL_PATH))

• 方法2、

保存整个模型，格式可以是pth.tar

 torch.save(model, PATH)

然后加载

model = torch.load(PATH)

还有一些在GPU和CPU之间的传输方法，具体见博客或者官方文档