读写Tensor

我们可以把Tensor存储到文件当中：

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

在同级目录下可以看到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.])}

读写模型

state_dict

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()
print(net.state_dict())
optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
print(optimizer.state_dict())
结果：
OrderedDict([('hidden.weight', tensor([[ 0.4458, -0.2268,  0.2473],
        [-0.3065,  0.0547,  0.0892]])), ('hidden.bias', tensor([0.0920, 0.5045])), ('output.weight', tensor([[0.0630, 0.2834]])), ('output.bias', tensor([0.5416]))])
{'state': {}, 'param_groups': [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [0, 1, 2, 3]}]}

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

保存与加载模型

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

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

保存和加载state_dict(推荐)

# 保存
torch.save(model.state_dict(), PATH) # 推荐的文件后缀名是pt或pth
# 加载
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))

特别要注意的是，在加载模型文件前先要对欲加载的模型进行实例化，然后再通过对象.load_state_dict(dict)方法载入模型，该方法不能直接传入文件名，需要先通过torch.load(path)方法导入字典。

例：

# 保存
torch.save(net.state_dict(), 'test.pt')
model = MLP()
# 载入
model.load_state_dict(torch.load('test.pt'))
print(model)
结果:
MLP(
  (hidden): Linear(in_features=3, out_features=2, bias=True)
  (act): ReLU()
  (output): Linear(in_features=2, out_features=1, bias=True)
)

要注意，要载入哪一个模型就尽量选择哪个模型类进行实例化，但结构完全相同的两个不同的网络类可以共享模型文件，参数或网络结构不一样的两个网络则不行。

如:

class MLP1(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)
# 保存
torch.save(net.state_dict(), 'test.pt')
model = MLP1()
# 载入
model.load_state_dict(torch.load('test.pt'))
print(model)
结果:
MLP1(
  (hidden): Linear(in_features=3, out_features=2, bias=True)
  (act): ReLU()
  (output): Linear(in_features=2, out_features=1, bias=True)
)

保存和加载整个模型

# 保存
torch.save(model, PATH)
# 加载
model = torch.load(PATH)

例如我们可以看看用MLP类实例化的对象运算得到的结果与从模型文件载入生成的对象得到的结果是否相等：

x = torch.rand(2, 3)
torch.save(net, 'MLP.pt')
model = torch.load('MLP.pt')
print(model(x) == net(x))
结果：
tensor([[True],
        [True]])

因为这net和net2都有同样的模型参数，那么对同一个输入X的计算结果将会是一样的。上面的输出也验证了这一点。此外，还有一些其他使用场景，例如GPU与CPU之间的模型保存与读取、使用多块GPU的模型的存储等等，使用的时候可以参考官方文档。

机器学习