读取和存储

《动手学深度学习pytorch》部分学习笔记,仅用作自己复习。

读取和存储

在实际中,我们有时需要把训练好的模型部署到很多不同的设备。在这种情况下,我们可以把内存中训练好的模型参数
存储在硬盘上供后续读取使用。 

 读写 TENSOR

可以直接使用 save 函数和 load 函数分别存储和读取 Tensor save 使⽤Pythonpickle实用程序将对象进行序列化,然后将序列列化的对象保存到disk,使⽤ save 可以保存各种对象,包括模型、张量和字典等。而 laod 使用pickle unpickle⼯具将pickle的对象⽂文件反序列化为内存。

下面的例子创建了 Tensor 变量 x ,并将其存在文件名同为 x.pt 的⽂件⾥。

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.])} 

读写模型

state_dict

PyTorch 中 , Module 的 可 学 习 参 数 ( 即 权 重 和 偏 差 ) , 模 块 模 型 包 含 在 参 数 中 ( 通 过model.parameters() 访问)。 state_dict 是⼀个从参数名称隐射到参数 Tesnor 的字典对象。

class MLP(nn.Module):
    def __init__(self):
        # 初始化
        super(MLP, self).__init__()
        # 隐藏层
        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()
# state_dict 是⼀个从参数名称隐射到参数 Tesnor 的字典对象。
net.state_dict()

输出:

OrderedDict([('hidden.weight', tensor([[ 0.2448, 0.1856, -0.5678],
[ 0.2030, -0.2073, -0.0104]])),
('hidden.bias', tensor([-0.3117, -0.4232])),
('output.weight', tensor([[-0.4556, 0.4084]])),
('output.bias', tensor([-0.3573]))])

注意,只有具有可学习参数的层(卷积层、线性层等)才有 state_dict 中的条目。优化器( optim )也有⼀个 state_dict ,其中包含关于优化器状态以及所使用的超参数的信息。

optimizer = torch.optim.SGD(net.parameters(), lr=0.001,
momentum=0.9)
optimizer.state_dict()

输出:{'param_groups': [{'dampening': 0,
'lr': 0.001,
'momentum': 0.9,
'nesterov': False,
'params': [4736167728, 4736166648, 4736167368, 4736165352],
'weight_decay': 0}],
'state': {}}

保存和加载模型

PyTorch中保存和加载训练模型有两种常见的⽅法:
1. 仅保存和加载模型参数( state_dict );
2. 保存和加载整个模型。

1. 保存和加载 state_dict (推荐方式)

保存:

torch.save(model.state_dict(), PATH) # 推荐的⽂件后缀名是pt或pth 

加载:

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))

2. 保存和加载整个模型

保存:

torch.save(model, PATH) 

加载:

model = torch.load(PATH) 

我们采⽤推荐的⽅法一来实验一下:

X = torch.randn(2, 3)
Y = net(X)
PATH = "./net.pt"
# 保存
torch.save(net.state_dict(), PATH)
net2 = MLP()
# 加载
net2.load_state_dict(torch.load(PATH))
Y2 = net2(X)
Y2 == Y

输出:

tensor([[1],
[1]], dtype=torch.uint8)

小结

  • 通过 save 函数和 load 函数可以很方便地读写 Tensor 。
  • 通过 save 函数和 load_state_dict 函数可以很⽅便地读写模型的参数。

猜你喜欢

转载自blog.csdn.net/dujuancao11/article/details/108468554