当前位置:   article > 正文

【深度学习】读取和存储训练好的模型参数(pyTorch)_训练结束后得到了一个pytorchmodel.bin文件,怎么去读取运用他

训练结束后得到了一个pytorchmodel.bin文件,怎么去读取运用他

读取、存储和部署模型

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')
  • 1
  • 2
  • 3
  • 4
  • 5
x2 = torch.load('x.pt')
  • 1
x2
  • 1
tensor([1., 1., 1.])
  • 1
# 存储一个Tensor列表并读回内存
y = torch.zeros(4)
torch.save([x ,y], 'xy.pt')
xy_list = torch.load('xy.pt')
xy_list
  • 1
  • 2
  • 3
  • 4
  • 5
[tensor([1., 1., 1.]), tensor([0., 0., 0., 0.])]
  • 1
# 存储并读取一个从字符串映射到Tensor的字典
torch.save({'x':x, 'y':y}, 'xy_dict.pt')
xy = torch.load('xy_dict.pt')
xy
  • 1
  • 2
  • 3
  • 4
{'x': tensor([1., 1., 1.]), 'y': tensor([0., 0., 0., 0.])}
  • 1

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()
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
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]))])
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
# 只有具有可学习参数的层(卷积层、线性层等)才有state_dict中的条目;
# 优化器(optim)也有一个state_dict,其中包含关于优化器状态以及所使用的超参数的信息。
optimizer =torch.optim.SGD(net.parameters(), lr = 0.001, momentum = 0.9)
optimizer.state_dict()

  • 1
  • 2
  • 3
  • 4
  • 5
{'state': {},
 'param_groups': [{'lr': 0.001,
   'momentum': 0.9,
   'dampening': 0,
   'weight_decay': 0,
   'nesterov': False,
   'params': [2581793307368, 2581793307448, 2581793307608, 2581793307688]}]}
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7

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
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
tensor([[-0.2372],
        [-0.1132]], grad_fn=<AddmmBackward>)





tensor([[-0.2372],
        [-0.1132]], grad_fn=<AddmmBackward>)
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9

参考原文

欢迎关注【OAOA

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/繁依Fanyi0/article/detail/772662
推荐阅读
相关标签
  

闽ICP备14008679号