动手学PyTorch | (17) 读取和存储

到目前为止，我们介绍了如何处理数据以及如何构建、训练和测试深度学习模型。然⽽在实际中，我们有时需要把训练好的模型部署到很多不同的设备。在这种情况下，我们可以把内存中训练好的模型参数存储在硬盘上供后续读取使用。

import torch
from torch import nn

目录

1. 读取Tensor

2. 读写模型

3. 小结

---

1. 读取Tensor

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

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

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

然后我们将数据从存储的⽂件读回内存。

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

我们还可以存储⼀个Tensor列表并读回内存：

y = torch.zeros(4)
torch.save([x, y], 'xy.pt')
xy_list = torch.load('xy.pt')
xy_list

存储并读取一个从字符串映射到Tensor的字典：

torch.save({'x': x, 'y': y}, 'xy_dict.pt')
xy = torch.load('xy_dict.pt')
xy

 

2. 读写模型

• state_dict

在PyTorch中，Module的可学习参数（权重和偏差），模块模型包含在参数中(通过model.parameter()访问)。state_dict是一个从参数名映射参数Tensor的字典对象。

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()
net.state_dict()

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

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

•  保存和加载模型

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

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

 因为net和net2有同样的参数，那么对同一个输入X的计算结果是一样的。

此外，还有⼀些其他使用场景，例如GPU与CPU之间的模型保存与读取、使⽤多块GPU的模型的存储等等，使⽤的时候可以参考官方中文文档。

3. 小结

1）通过save函数和load函数可以方便的读写Tensor。

2）通过save函数和load_state_dict函数可以很方便地读写模型的参数。