【深度学习】读取和存储训练好的模型参数（pyTorch）

读取、存储和部署模型

1 读写Tensor

可以直接使用save函数和load函数分别存储和读取Tensor。

save使用Python的pickle实用程序将对象进行序列化，然后将序列化的对象保存到disk，使用save可以保存各种对象,包括模型、张量和字典等;

而load使用pickle unpickle工具将pickle的对象文件反序列化为内存。

import torch 
from torch import nn

x=torch.ones(3)
torch.save(x, 'x.pt')

x2 = torch.load('x.pt')

x2

tensor([1., 1., 1.])

# 存储一个Tensor列表并读回内存
y = torch.zeros(4)
torch.save([x ,y], 'xy.pt')
xy_list = torch.load('xy.pt')
xy_list

[tensor([1., 1., 1.]), tensor([0., 0., 0., 0.])]

# 存储并读取一个从字符串映射到Tensor的字典
torch.save({'x':x, 'y':y}, 'xy_dict.pt')
xy = torch.load('xy_dict.pt')
xy

{'x': tensor([1., 1., 1.]), 'y': tensor([0., 0., 0., 0.])}

2 读写模型

# state_dict是一个从参数名称隐射到参数Tesnor的字典对象。
class MLP(nn.Module):
    def __init__(self, **kwargs): # **kwargs表示如果我们不知道要往函数中传入多少个关键词参数，或者想传入字典的值作为关键词参数时，那就要使用**kwargs。
        super(MLP, self).__init__(**kwargs)
        self.hidden = nn.Linear(3, 2)
        self.act = nn.ReLU()
        self.output = nn.Linear(2, 1)
    def forward(self, x):
        a = self.act(self.hidden(x))
        return self.output(a)
net = MLP()
net.state_dict()

OrderedDict([('hidden.weight',
              tensor([[-0.4192,  0.1470,  0.2054],
                      [ 0.1476, -0.2715,  0.1462]])),
             ('hidden.bias', tensor([-0.2444, -0.2241])),
             ('output.weight', tensor([[ 0.3270, -0.3903]])),
             ('output.bias', tensor([-0.1880]))])

# 只有具有可学习参数的层(卷积层、线性层等)才有state_dict中的条目；
# 优化器(optim)也有一个state_dict，其中包含关于优化器状态以及所使用的超参数的信息。
optimizer =torch.optim.SGD(net.parameters(), lr = 0.001, momentum = 0.9)
optimizer.state_dict()

{'state': {},
 'param_groups': [{'lr': 0.001,
   'momentum': 0.9,
   'dampening': 0,
   'weight_decay': 0,
   'nesterov': False,
   'params': [2581793307368, 2581793307448, 2581793307608, 2581793307688]}]}

PyTorch中保存和加载训练模型有两种常见的方法:

1. 仅保存和加载模型参数(state_dict)；
2. 保存和加载整个模型

# 保存和加载模型参数(state_dict)

# 保存
X = torch.randn(2, 3)
Y = net(X)
print(Y)
path = './net.pt'
torch.save(net.state_dict(), path)
net2 = MLP()
net2.load_state_dict(torch.load(path))
Y2 = net2(X)
Y2

tensor([[-0.2372],
        [-0.1132]], grad_fn=<AddmmBackward>)





tensor([[-0.2372],
        [-0.1132]], grad_fn=<AddmmBackward>)

参考原文

欢迎关注【OAOA】