Pytorch保存训练模型

简介

我们在使用Pytorch深度学习框架训练自己的模型时，需要保存训练结果以便于下次直接使用训练好的模型对新数据进行预测。本文将介绍Pytorch官方提供的两种模型保存方法

内容

方法1：仅保存模型参数

代码如下：

# 保存模型状态字典，PATH为保存的路径以及文件名
# 例如： PATH='models/myModel.pt'
torch.save(model.state_dict(), PATH)

# 加载模型
# 先实例化一个模型
model = TheModelClass(...)
# 加载状态字典
model.load_state_dict(torch.load(PATH))

方法2：保存整个模型

# 保存
torch.save(model, PATH)
# 加载
model = torch.load(PATH)