torch模型保存

torch模型保存与加载

torch保存与记载核心就两个方法，两个方式：

1. torch.save() - 作用是将对象保存到序列化文件中。

   • 方式一：保存模型 - torch.save(model, ‘./model.pth’) 或者保存一个Python对象- torch.save({‘save’:‘ddd’}, ‘./adict.pth’)
   • 方式二：保存模型参数 - torch.save(model.state_dict(), ‘./model.pth’)

2. 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 optim


class 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 output
    

 model = Model(12, 12)
# 方式一
torch.save(model, './model.pth')

# 方式二
torch.save(model.state_dict(), './model_state_dict.pth')

2. 加载模型

• 加载方式一，注意当前文件中要有Model这个类，但不用实例化。

from models import Model

model = torch.load('./model.pth')

• 加载方式二，注意要实例化模型这个类，才能调用load_state_dict()这个方法。
• 方式二的本质就是返回一个字典。

from models import Model

model = Model()
model.load_state_dict(torch.load('./model_state_dict.pth'))