[pytorch、学习] - 4.5 读取和存储

参考

4.5 读取和存储

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

4.5.1 读写tensor

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

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

import torch
import torch.nn as nn

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

4.5.2 读写模型

4.5.2.1 state_dict

static_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中的条目

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

4.5.2.2 保存和加载模型

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

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

1. 保存加载static_dict(推荐方式)
torch.save(model.state_dict(), PATH)

# 保存
torch.save(model.state_dict(), PATH)

# 加载
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.laod(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