Pytorch数据的读取与存储

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

一、读写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')
print(x2)

输出：

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

我们还可以存储一个Tensor列表并读回内存。

y=torch.zeros(4)
torch.save([x,y],'xy.pt')
xy_list=torch.load('xy.pt')
print(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')
print(xy)

输出：

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

二、读写模型

2.1 state_dict

Pytorch中，Module的可学习参数（即权重和偏差）以及模块模型包含在参数中（通过model.parameters()访问）。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()
print(net.state_dict())

输出：

OrderedDict([('hidden.weight', tensor([[-0.3099, -0.3298,  0.2657],
        [-0.3176,  0.3805,  0.0322]])), ('hidden.bias', tensor([-0.0345, -0.0435])), ('output.weight', tensor([[-0.1456,  0.0254]])), ('output.bias', tensor([0.3656]))])

注意，只有具有可学习参数的层（卷积层、线性层等）才有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': [0, 1, 2, 3]}]}

2.2 保存和加载模型

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

1. 仅保存和加载模型参数（state_dict）：
2. 保存和加载整个模型。

2.2.1 保存和加载state_dict（推荐方式）

保存：

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

加载：

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

2.2.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)
print(Y2)
print(Y)

输出：

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

因为net和net2具有同样的模型参数，所以对同一个输入X的计算结果是一样的，上面的输出也证明了这点。