torch模型保存与加载
torch保存与记载核心就两个方法，两个方式：
torch.save() - 作用是将对象保存到序列化文件中。
- 方式一：保存模型 - torch.save(model, ‘./model.pth’) 或者保存一个Python对象- torch.save({‘save’:‘ddd’}, ‘./adict.pth’)
- 方式二：保存模型参数 - torch.save(model.state_dict(), ‘./model.pth’)
torch.load() - 作用是将序列化文件生成回对象。
加载方式一的模型： model = torch.load(‘./model.pth’)
加载方式二的参数:
model = Model()
model.load_state_dict(torch.load(‘./model.pth’))
1. 保存模型
- torch.save可以将任何对象保存到序列话文件中，并不一定是torch对象。
- 方式二中torch.save就是保存的字典。
import torch
import torch.nn as nn
import torch.optim as optimclass Model(nn.Module):def __init__(self, input_size, output_size):super(Model, self).__init__()self.linear1 = nn.Linear(input_size, input_size * 2)self.linear2 = nn.Linear(input_size * 2, output_size)def forward(self, inputs):inputs = self.linear1(inputs)output = self.linear2(inputs)return outputmodel = Model(12, 12)
# 方式一
torch.save(model, './model.pth')# 方式二
torch.save(model.state_dict(), './model_state_dict.pth')
2. 加载模型
- 加载方式一，注意当前文件中要有Model这个类，但不用实例化。
from models import Modelmodel = torch.load('./model.pth')
- 加载方式二，注意要实例化模型这个类，才能调用load_state_dict()这个方法。
- 方式二的本质就是返回一个字典。
from models import Modelmodel = Model()
model.load_state_dict(torch.load('./model_state_dict.pth'))
)
