两种保存方式:
- 保存模型和参数,对应加载时直接赋给model
- 仅仅只保存参数,对应加载时需要先实例化一个model,再把参数加载给model
第一种:
import torch
#保存
torch.save(model,"./weight/best_model.pth")
#加载
best_model = torch.load("./weight/best_model.pth")
第二种
import torch
#仅保存参数
torch.save(model.state_dict(), './weight/best_model.pth')
#加载
model = MyModel() # 实例化
state_dict = torch.load('model_state_dict.pth')#加载参数
model.load_state_dict(state_dict)#将参数加载给model